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)