Load train dataset in SegNet

I’m using this Pytorch implementation of Segnet with pretrained values I found for object segmentation, and it works fine. Now I want to resume the training from the values I have, using a new dataset with similar images. How can I do that?

“train.py” doesn’t have a method to load the data and I don’t know how to replace the commented part in order to make the code working. This is my first time posting and using Pytorch, I’m sorry if I’m not formatting right or if the question is a bit dull.

“train” method found in “train.py”:

def train(epoch):

model.train()
# update learning rate
lr = args.lr * (0.1 ** (epoch // 30))
for param_group in optimizer.param_groups:
    param_group['lr'] = lr

# define a weighted loss (0 weight for 0 label)
weights_list = [0]+[1 for i in range(17)]
weights = np.asarray(weights_list)
weigthtorch = torch.Tensor(weights_list)
if(USE_CUDA):
    loss = nn.CrossEntropyLoss(weight=weigthtorch).cuda()
else:
    loss = nn.CrossEntropyLoss(weight=weigthtorch)


total_loss = 0

# iteration over the batches
batches = []
for batch_idx,batch_files in enumerate(tqdm(batches)):

    # containers
    batch = np.zeros((args.batch_size,input_nbr, imsize, imsize), dtype=float)
    batch_labels = np.zeros((args.batch_size,imsize, imsize), dtype=int)

    # fill the batch
    # ... I think I should insert some black magic here

    batch_th = Variable(torch.Tensor(batch))
    target_th = Variable(torch.LongTensor(batch_labels))

    if USE_CUDA:
        batch_th =batch_th.cuda()
        target_th = target_th.cuda()

    # initilize gradients
    optimizer.zero_grad()

    # predictions
    output = model(batch_th)

    # Loss
    output = output.view(output.size(0),output.size(1), -1)
    output = torch.transpose(output,1,2).contiguous()
    output = output.view(-1,output.size(2))
    target = target.view(-1)

    l_ = loss(output.cuda(), target)
    total_loss += l_.cpu().data.numpy()
    l_.cuda()
    l_.backward()
    optimizer.step()

return total_loss/len(files)