def pretrain(dataloader, network, path=None): class config: lr = 1e-3 epochs = 100 path = '../checkpoint/pretrained_net.pth' pretrain_config = config() if path: pretrain_config.path = path network.to(device) criterion_ = CrossEntropyLoss2d() optimiser_ = torch.optim.Adam(network.parameters(), pretrain_config.lr) loss_meter = AverageValueMeter() for i in range(pretrain_config.epochs): loss_meter.reset() for i, (img, mask, weak_mask, _) in tqdm(enumerate(dataloader)): img, mask = img.to(device), mask.to(device) optimiser_.zero_grad() output = network(img) loss = criterion_(output, mask.squeeze(1)) loss.backward() optimiser_.step() loss_meter.add(loss.item()) # import ipdb # ipdb.set_trace() print(loss_meter.value()[0]) torch.save(network.state_dict(), pretrain_config.path) # torch.save(network.parameters(),path) print('pretrained model saved.')
def __init__(self, neural_network, lowerbound, upperbound): super(networks, self).__init__() self.lowbound = lowerbound self.upbound = upperbound self.neural_net = neural_network self.reset() self.optimiser = torch.optim.Adam(self.neural_net.parameters(), lr=0.001, weight_decay=1e-5) self.CEloss_criterion = CrossEntropyLoss2d() self.u_r = 1 self.u_s = 1
def __init__(self, neural_network, lowerbound, upperbound): super(networks, self).__init__() self.lowbound = lowerbound self.upbound = upperbound self.neural_net = neural_network self.reset() self.optimiser = torch.optim.Adam(self.neural_net.parameters(), lr=0.0001) self.CEloss_criterion = CrossEntropyLoss2d() self.u_r = 1000 self.lamda = 0 # self.set_bound=False self.sigma = 0.05 self.scale = 0.5
def update_theta(f_theta_labeled, f_theta_unlabeled, gamma, s, u, v): global u_r, u_s, net global labeled_img, labeled_mask, labeled_weak_mask, unlabeled_img, unlabeled_mask optimiser = torch.optim.Adam(net.parameters(), lr=1e-3) criterion = CrossEntropyLoss2d() for i in xrange(10): loss_labeled = criterion(f_theta_labeled, labeled_mask.squeeze(1)) loss_unlabeled = u_r * (f_theta_unlabeled - gamma.float() + torch.Tensor(u).float()).norm(p=2) + \ u_s * (f_theta_unlabeled - s.float() + torch.Tensor(v)).norm(p=2) loss = loss_labeled + loss_unlabeled optimiser.zero_grad() loss.backward() optimiser.step() f_theta_labeled = net(labeled_img) # shape b,c,w,h f_theta_unlabeled = net(unlabeled_img) return f_theta_labeled, f_theta_unlabeled
def train(self): self.color_transform = Colorization(2) # Dataset loader for train and test dataset_train = DataLoader(PCD( os.path.join(self.args.datadir, 'set{}'.format(self.args.cvset), 'train')), num_workers=self.args.num_workers, batch_size=self.args.batch_size, shuffle=True) self.dataset_test = PCD( os.path.join(self.args.datadir, 'set{}'.format(self.args.cvset), 'test')) self.test_path = os.path.join(self.dn_save, 'test') if not os.path.exists(self.test_path): os.makedirs(self.test_path) # Set loss function, optimizer and learning rate weight = torch.ones(2) criterion = CrossEntropyLoss2d(weight.cuda()) optimizer = Adam(self.model.parameters(), lr=0.0001, betas=(0.5, 0.999)) lambda1 = lambda icount: (float)(self.args.max_iteration - icount) / ( float)(self.args.max_iteration) model_lr_scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda1) fn_loss = os.path.join(self.dn_save, 'loss.csv') f_loss = open(fn_loss, 'w') writer = csv.writer(f_loss) self.writers = SummaryWriter(os.path.join(self.dn_save, 'log')) # Training loop icount_loss = [] while self.icount < self.args.max_iteration: for step, (inputs_train, mask_train) in enumerate(dataset_train): inputs_train = inputs_train.cuda() mask_train = mask_train.cuda() inputs_train = Variable(inputs_train) mask_train = Variable(mask_train) outputs_train = self.model(inputs_train) optimizer.zero_grad() self.loss = criterion(outputs_train, mask_train[:, 0]) self.loss.backward() optimizer.step() self.icount += 1 icount_loss.append(self.loss.item()) writer.writerow([self.icount, self.loss.item()]) if self.args.icount_plot > 0 and self.icount % self.args.icount_plot == 0: self.test() average = sum(icount_loss) / len(icount_loss) print('loss: {0} (icount: {1})'.format( average, self.icount)) icount_loss.clear() if self.args.icount_save > 0 and self.icount % self.args.icount_save == 0: self.checkpoint() # Call lr_schduler.step() after optimizer.step() model_lr_scheduler.step() f_loss.close()
def train(args, model, enc=False): global best_acc #TODO: calculate weights by processing dataset histogram (now its being set by hand from the torch values) #create a loder to run all images and calculate histogram of labels, then create weight array using class balancing weight = torch.ones(NUM_CLASSES) if (enc): weight[0] = 2.3653597831726 weight[1] = 4.4237880706787 weight[2] = 2.9691488742828 weight[3] = 5.3442072868347 weight[4] = 5.2983593940735 weight[5] = 5.2275490760803 weight[6] = 5.4394111633301 weight[7] = 5.3659925460815 weight[8] = 3.4170460700989 weight[9] = 5.2414722442627 weight[10] = 4.7376127243042 weight[11] = 5.2286224365234 weight[12] = 5.455126285553 weight[13] = 4.3019247055054 weight[14] = 5.4264230728149 weight[15] = 5.4331531524658 weight[16] = 5.433765411377 weight[17] = 5.4631009101868 weight[18] = 5.3947434425354 else: weight[0] = 2.8149201869965 weight[1] = 6.9850029945374 weight[2] = 3.7890393733978 weight[3] = 9.9428062438965 weight[4] = 9.7702074050903 weight[5] = 9.5110931396484 weight[6] = 10.311357498169 weight[7] = 10.026463508606 weight[8] = 4.6323022842407 weight[9] = 9.5608062744141 weight[10] = 7.8698215484619 weight[11] = 9.5168733596802 weight[12] = 10.373730659485 weight[13] = 6.6616044044495 weight[14] = 10.260489463806 weight[15] = 10.287888526917 weight[16] = 10.289801597595 weight[17] = 10.405355453491 weight[18] = 10.138095855713 weight[19] = 0 #loader = DataLoader(VOC12(args.datadir, input_transform, target_transform), # num_workers=args.num_workers, batch_size=args.batch_size, shuffle=True) assert os.path.exists( args.datadir), "Error: datadir (dataset directory) could not be loaded" co_transform = MyCoTransform(enc, augment=True, height=args.height) #1024) co_transform_val = MyCoTransform(enc, augment=False, height=args.height) #1024) dataset_train = cityscapes(args.datadir, co_transform, 'train') dataset_val = cityscapes(args.datadir, co_transform_val, 'val') loader = DataLoader(dataset_train, num_workers=args.num_workers, batch_size=args.batch_size, shuffle=True) loader_val = DataLoader(dataset_val, num_workers=args.num_workers, batch_size=args.batch_size, shuffle=False) if args.cuda: weight = weight.cuda() criterion = CrossEntropyLoss2d(weight) print(type(criterion)) savedir = args.dirpath + args.savedir if (enc): automated_log_path = savedir + "/automated_log_encoder.txt" modeltxtpath = savedir + "/model_encoder.txt" else: automated_log_path = savedir + "/automated_log.txt" modeltxtpath = savedir + "/model.txt" if (not os.path.exists(automated_log_path) ): #dont add first line if it exists with open(automated_log_path, "a") as myfile: myfile.write( "Epoch\t\tTrain-loss\t\tTest-loss\t\tTrain-IoU\t\tTest-IoU\t\tlearningRate" ) with open(modeltxtpath, "w") as myfile: myfile.write(str(model)) #TODO: reduce memory in first gpu: https://discuss.pytorch.org/t/multi-gpu-training-memory-usage-in-balance/4163/4 #https://github.com/pytorch/pytorch/issues/1893 """ #Some optimizer examples: optimizer = Adam(model.parameters()) if args.model.startswith('FCN'): optimizer = SGD(model.parameters(), 1e-4, .9, 2e-5) if args.model.startswith('PSP'): optimizer = SGD(model.parameters(), 1e-2, .9, 1e-4) if args.model.startswith('Seg'): optimizer = SGD(model.parameters(), 1e-3, .9) if args.model.startswith('E'): #optimizer = Adam(model.parameters(), 1e-3, .9) optimizer = Adam(model.parameters(), 5e-4, .9, weight_decay=2e-4) #5e-4 wd: 2e-4 """ #optimizer = Adam(model.parameters(), 5e-4, (0.9, 0.999), eps=1e-08, weight_decay=2e-4) ## scheduler 1 optimizer = Adam(model.parameters(), 5e-4, (0.9, 0.999), eps=1e-08, weight_decay=1e-4) ## scheduler 2 start_epoch = 1 if args.resume: #Must load weights, optimizer, epoch and best value. if enc: filenameCheckpoint = savedir + '/checkpoint_enc.pth.tar' else: filenameCheckpoint = savedir + '/checkpoint.pth.tar' assert os.path.exists( filenameCheckpoint ), "Error: resume option was used but checkpoint was not found in folder" checkpoint = torch.load(filenameCheckpoint) start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) best_acc = checkpoint['best_acc'] print("=> Loaded checkpoint at epoch {})".format(checkpoint['epoch'])) #scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.5) # set up scheduler ## scheduler 1 lambda1 = lambda epoch: pow( (1 - ((epoch - 1) / args.num_epochs)), 0.9) ## scheduler 2 scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda1) ## scheduler 2 if args.visualize and args.steps_plot > 0: board = Dashboard(args.port) for epoch in range(start_epoch, args.num_epochs + 1): print("----- TRAINING - EPOCH", epoch, "-----") scheduler.step(epoch) ## scheduler 2 epoch_loss = [] time_train = [] doIouTrain = args.iouTrain doIouVal = args.iouVal if (doIouTrain): iouEvalTrain = iouEval(NUM_CLASSES) usedLr = 0 for param_group in optimizer.param_groups: print("LEARNING RATE: ", param_group['lr']) usedLr = float(param_group['lr']) model.train() for step, (images, labels) in enumerate(loader): print(step) start_time = time.time() #print (labels.size()) #print (np.unique(labels.numpy())) #print("labels: ", np.unique(labels[0].numpy())) #labels = torch.ones(4, 1, 512, 1024).long() if args.cuda: images = images.cuda() labels = labels.cuda() inputs = Variable(images) targets = Variable(labels) outputs = model(inputs, only_encode=enc) #print("targets", np.unique(targets[:, 0].cpu().data.numpy())) optimizer.zero_grad() loss = criterion(outputs, targets[:, 0]) loss.backward() optimizer.step() epoch_loss.append(loss.data[0]) time_train.append(time.time() - start_time) print(loss.data[0]) if (doIouTrain): #start_time_iou = time.time() iouEvalTrain.addBatch( outputs.max(1)[1].unsqueeze(1).data, targets.data) #print ("Time to add confusion matrix: ", time.time() - start_time_iou) #print(outputs.size()) if args.visualize and args.steps_plot > 0 and step % args.steps_plot == 0: start_time_plot = time.time() image = inputs[0].cpu().data #image[0] = image[0] * .229 + .485 #image[1] = image[1] * .224 + .456 #image[2] = image[2] * .225 + .406 #print("output", np.unique(outputs[0].cpu().max(0)[1].data.numpy())) board.image(image, 'input (epoch: {epoch}, step: {step})') if isinstance(outputs, list): #merge gpu tensors board.image( color_transform( outputs[0][0].cpu().max(0)[1].data.unsqueeze(0)), 'output (epoch: {epoch}, step: {step})') else: board.image( color_transform( outputs[0].cpu().max(0)[1].data.unsqueeze(0)), 'output (epoch: {epoch}, step: {step})') board.image(color_transform(targets[0].cpu().data), 'target (epoch: {epoch}, step: {step})') print("Time to paint images: ", time.time() - start_time_plot) if args.steps_loss > 0 and step % args.steps_loss == 0: average = sum(epoch_loss) / len(epoch_loss) print( 'loss: {average:0.4} (epoch: {epoch}, step: {step})', "// Avg time/img: %.4f s" % (sum(time_train) / len(time_train) / args.batch_size)) average_epoch_loss_train = sum(epoch_loss) / len(epoch_loss) iouTrain = 0 if (doIouTrain): iouTrain, iou_classes = iouEvalTrain.getIoU() iouStr = getColorEntry(iouTrain) + '{:0.2f}'.format( iouTrain * 100) + '\033[0m' print("EPOCH IoU on TRAIN set: ", iouStr, "%") #Validate on 500 val images after each epoch of training print("----- VALIDATING - EPOCH", epoch, "-----") model.eval() epoch_loss_val = [] time_val = [] if (doIouVal): iouEvalVal = iouEval(NUM_CLASSES) for step, (images, labels) in enumerate(loader_val): start_time = time.time() if args.cuda: images = images.cuda() labels = labels.cuda() inputs = Variable( images, volatile=True ) #volatile flag makes it free backward or outputs for eval targets = Variable(labels, volatile=True) outputs = model(inputs, only_encode=enc) loss = criterion(outputs, targets[:, 0]) epoch_loss_val.append(loss.data[0]) time_val.append(time.time() - start_time) #Add batch to calculate TP, FP and FN for iou estimation if (doIouVal): #start_time_iou = time.time() iouEvalVal.addBatch( outputs.max(1)[1].unsqueeze(1).data, targets.data) #print ("Time to add confusion matrix: ", time.time() - start_time_iou) if args.visualize and args.steps_plot > 0 and step % args.steps_plot == 0: start_time_plot = time.time() image = inputs[0].cpu().data board.image(image, 'VAL input (epoch: {epoch}, step: {step})') if isinstance(outputs, list): #merge gpu tensors board.image( color_transform( outputs[0][0].cpu().max(0)[1].data.unsqueeze(0)), 'VAL output (epoch: {epoch}, step: {step})') else: board.image( color_transform( outputs[0].cpu().max(0)[1].data.unsqueeze(0)), 'VAL output (epoch: {epoch}, step: {step})') board.image(color_transform(targets[0].cpu().data), 'VAL target (epoch: {epoch}, step: {step})') print("Time to paint images: ", time.time() - start_time_plot) if args.steps_loss > 0 and step % args.steps_loss == 0: average = sum(epoch_loss_val) / len(epoch_loss_val) print( 'VAL loss: {average:0.4} (epoch: {epoch}, step: {step})', "// Avg time/img: %.4f s" % (sum(time_val) / len(time_val) / args.batch_size)) average_epoch_loss_val = sum(epoch_loss_val) / len(epoch_loss_val) #scheduler.step(average_epoch_loss_val, epoch) ## scheduler 1 # update lr if needed iouVal = 0 if (doIouVal): iouVal, iou_classes = iouEvalVal.getIoU() iouStr = getColorEntry(iouVal) + '{:0.2f}'.format( iouVal * 100) + '\033[0m' print("EPOCH IoU on VAL set: ", iouStr, "%") # remember best valIoU and save checkpoint if iouVal == 0: current_acc = average_epoch_loss_val else: current_acc = iouVal is_best = current_acc > best_acc best_acc = max(current_acc, best_acc) if enc: filenameCheckpoint = savedir + '/checkpoint_enc.pth.tar' filenameBest = savedir + '/model_best_enc.pth.tar' else: filenameCheckpoint = savedir + '/checkpoint.pth.tar' filenameBest = savedir + '/model_best.pth.tar' save_checkpoint( { 'epoch': epoch + 1, 'arch': str(model), 'state_dict': model.state_dict(), 'best_acc': best_acc, 'optimizer': optimizer.state_dict(), }, is_best, filenameCheckpoint, filenameBest) #SAVE MODEL AFTER EPOCH if (enc): filename = savedir + '/model_encoder-' + '{epoch:03}.pth' filenamebest = savedir + '/model_encoder_best.pth' else: filename = savedir + '/model-{epoch:03}.pth' filenamebest = savedir + '/model_best.pth' if args.epochs_save > 0 and step > 0 and step % args.epochs_save == 0: torch.save(model.state_dict(), filename) print('save: {filename} (epoch: {epoch})') if (is_best): torch.save(model.state_dict(), filenamebest) print('save: {filenamebest} (epoch: {epoch})') if (not enc): with open(savedir + "/best.txt", "w") as myfile: myfile.write("Best epoch is %d, with Val-IoU= %.4f" % (epoch, iouVal)) else: with open(savedir + "/best_encoder.txt", "w") as myfile: myfile.write("Best epoch is %d, with Val-IoU= %.4f" % (epoch, iouVal)) #SAVE TO FILE A ROW WITH THE EPOCH RESULT (train loss, val loss, train IoU, val IoU) #Epoch Train-loss Test-loss Train-IoU Test-IoU learningRate with open(automated_log_path, "a") as myfile: myfile.write("\n%d\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.8f" % (epoch, average_epoch_loss_train, average_epoch_loss_val, iouTrain, iouVal, usedLr)) return (model) #return model (convenience for encoder-decoder training)
def train(args, model): model.train() weight = torch.ones(NUM_CLASSES) #weight[0] = 0 loader = DataLoader(CityScapes(args.datadir, input_transform, target_transform), num_workers=args.num_workers, batch_size=args.batch_size, shuffle=True) val_loader = DataLoader(CityScapes_validation(args.datadir, input_transform, target_transform), num_workers=args.num_workers, batch_size=args.batch_size, shuffle=True) if args.cuda: criterion = CrossEntropyLoss2d(weight.cuda()) #criterion=torch.nn.BCEWithLogitsLoss() else: criterion = CrossEntropyLoss2d(weight) optimizer = SGD(model.parameters(), 1e-4, .9, 2e-5) ''' if args.model.startswith('FCN'): optimizer = SGD(model.parameters(), 1e-4, .9, 2e-5) if args.model.startswith('PSP'): optimizer=SGD(filter(lambda p: p.requires_grad, model.parameters()), 1e-2,0.9,1e-4) #optimizer = SGD(model.parameters(), 1e-2, .9, 1e-4) if args.model.startswith('Seg'): optimizer = SGD(filter(lambda p: p.requires_grad, model.parameters()), 1e-3, .9) ''' print("Total steps:",len(loader)) best_loss=100 best_val_loss=100 best_acc=0 best_val_acc=0 for epoch in range(1, args.num_epochs+1): epoch_loss = [] iteration=1 train_acc=[] for step, (images, labels) in enumerate(loader): print("Iter:"+str(iteration)) iteration=iteration+1 if args.cuda: images = images.cuda() labels = labels.cuda() inputs = Variable(images) targets = Variable(labels) outputs=model(inputs) optimizer.zero_grad() loss = criterion(outputs, targets[:, 0,:,:])#Bx1xHxW loss.backward() optimizer.step() print(loss.item()) epoch_loss.append(loss.item()) acc_vec=[] for b in range(outputs.size()[0]): acc_vec.append(pixel_accuracy(torch.max(outputs[b,:,:,:],0)[1],targets[b,0,:,:])) acc=sum(acc_vec)/len(acc_vec) print("train_acc: "+str(acc)) train_acc.append(acc) if args.steps_loss > 0 and step>0 and step % args.steps_loss == 0: average = sum(epoch_loss) / len(epoch_loss) average_acc=sum(train_acc) / len(train_acc) epoch_loss=[] train_acc=[] if best_loss>average: best_loss=average #torch.save(model.state_dict(), "model_linknet34.pth") #print("Model saved!") if best_acc<average_acc: best_acc=average_acc print("loss: "+str(average)+" epoch: "+str(epoch)+", step: "+str(step)) print("best loss: "+str(best_loss)+" epoch: "+str(epoch)+", step: "+str(step)) print("train acc: "+str(average_acc)) print("best train acc: "+str(best_acc)) f=open("train_loss.txt","a") f.write(str(epoch)+" "+str(step)+" "+str(average)+" "+str(best_loss)+" "+str(average_acc)+" "+str(best_acc)+"\n") f.close() print("Best loss: "+str(best_loss)) print("Best val loss: "+str(best_val_loss)) print("best train acc: "+str(best_acc)) print("Best val acc: "+str(best_val_acc)) epoch_loss = [] val_acc=[] iteration=1 for step, (images, labels) in enumerate(val_loader): print("Val Iter:"+str(iteration)) iteration=iteration+1 if args.cuda: images = images.cuda() labels = labels.cuda() inputs = Variable(images) targets = Variable(labels) outputs=model(inputs) loss = criterion(outputs, targets[:, 0,:,:]) print(loss.item()) epoch_loss.append(loss.item()) val_acc_vec=[] for b in range(outputs.size()[0]): val_acc_vec.append(pixel_accuracy(torch.max(outputs[b,:,:,:],0)[1],targets[b,0,:,:])) acc=sum(val_acc_vec)/len(val_acc_vec) val_acc.append(acc) if args.steps_loss > 0 and step>0 and step % args.steps_loss == 0: average = sum(epoch_loss) / len(epoch_loss) average_acc=sum(val_acc) / len(val_acc) epoch_loss=[] val_acc=[] if best_val_loss>average: best_val_loss=average torch.save(model.state_dict(), "model_linknet34.pth") print("Model saved!") if best_val_acc<average_acc: best_val_acc=average_acc print("val loss: "+str(average)+" epoch: "+str(epoch)+", step: "+str(step)) print("best val loss: "+str(best_val_loss)+" epoch: "+str(epoch)+", step: "+str(step)) print("val acc: "+str(average_acc)) print("best val acc: "+str(best_acc)) f1=open("val_loss.txt","a") f1.write(str(epoch)+" "+str(step)+" "+str(average)+" "+str(best_val_loss)+" "+ str(average_acc)+" "+str(best_val_acc)+"\n") f1.close() print("Best val loss: "+str(best_val_loss)) print("Best val acc: "+str(best_val_acc))
loader = data.DataLoader(VOCDataSet( "/home/zeng/data/datasets/segmentation_Dataset", img_transform=input_transform, label_transform=target_transform), batch_size=12, shuffle=True, pin_memory=True) res101 = resnet101(pretrained=True).cuda() seg = Seg().cuda() weight = torch.ones(22) weight[21] = 0 criterion = CrossEntropyLoss2d(weight.cuda()) optimizer_seg = torch.optim.Adam(seg.parameters(), lr=1e-3) optimizer_feat = torch.optim.Adam(res101.parameters(), lr=1e-4) for t in range(10): for i, (img, label) in enumerate(loader): img = img.cuda() label = label[0].cuda() label = Variable(label) input = Variable(img) feats = res101(input) output = seg(feats) seg.zero_grad()
feature.load_state_dict(torch.load(pretrained_feature_file)) deconv = Deconv(opt.i) deconv.cuda() if resume_ep >= 0: feature_param_file = glob.glob('%s/feature-epoch-%d*.pth'%(check_dir, resume_ep)) deconv_param_file = glob.glob('%s/deconv-epoch-%d*.pth'%(check_dir, resume_ep)) feature.load_state_dict(torch.load(feature_param_file[0])) deconv.load_state_dict(torch.load(deconv_param_file[0])) train_loader = torch.utils.data.DataLoader( MyBoxPixData(train_dir, transform=True, crop=True, hflip=True, vflip=False, source=opt.q), batch_size=bsize, shuffle=True, num_workers=4, pin_memory=True) criterion = CrossEntropyLoss2d(weight=torch.FloatTensor(label_weight)) criterion.cuda() optimizer_deconv = torch.optim.Adam(deconv.parameters(), lr=1e-3) optimizer_feature = torch.optim.Adam(feature.parameters(), lr=1e-4) for it in range(resume_ep+1, iter_num): for ib, (data, lbl) in enumerate(train_loader): inputs = Variable(data).cuda() lbl = Variable(lbl.long()).cuda() feats = feature(inputs) msk = deconv(feats) msk = functional.upsample(msk, scale_factor=8) loss = criterion(msk, lbl)
def train(args, model): model.train() weight = torch.ones(22) weight[0] = 0 loader = DataLoader(VOC12(args.datadir, input_transform, target_transform), num_workers=args.num_workers, batch_size=args.batch_size, shuffle=True) if args.cuda: criterion = CrossEntropyLoss2d(weight.cuda()) else: criterion = CrossEntropyLoss2d(weight) optimizer = Adam(model.parameters()) if args.model.startswith('FCN'): optimizer = SGD(model.parameters(), 1e-4, .9, 2e-5) if args.model.startswith('PSP'): optimizer = SGD(model.parameters(), 1e-2, .9, 1e-4) if args.model.startswith('Seg'): optimizer = SGD(model.parameters(), 1e-3, .9) if args.steps_plot > 0: board = Dashboard(args.port) for epoch in range(1, args.num_epochs + 1): epoch_loss = [] for step, (images, labels) in enumerate(loader): if args.cuda: images = images.cuda() labels = labels.cuda() inputs = Variable(images) targets = Variable(labels) outputs = model(inputs) optimizer.zero_grad() loss = criterion(outputs, targets[:, 0]) loss.backward() optimizer.step() epoch_loss.append(loss.data[0]) if args.steps_plot > 0 and step % args.steps_plot == 0: image = inputs[0].cpu().data image[0] = image[0] * .229 + .485 image[1] = image[1] * .224 + .456 image[2] = image[2] * .225 + .406 board.image(image, f('input (epoch: {epoch}, step: {step})')) board.image(color_transform(outputs[0].cpu().max(0)[1].data), f('output (epoch: {epoch}, step: {step})')) board.image(color_transform(targets[0].cpu().data), f('target (epoch: {epoch}, step: {step})')) if args.steps_loss > 0 and step % args.steps_loss == 0: average = sum(epoch_loss) / len(epoch_loss) print(f('loss: {average} (epoch: {epoch}, step: {step})')) if args.steps_save > 0 and step % args.steps_save == 0: filename = f('{args.model}-{epoch:03}-{step:04}.pth') torch.save(model.state_dict(), filename) print(f('save: {filename} (epoch: {epoch}, step: {step})'))
def train(args, model): model.train() weight = torch.ones(22) weight[0] = 0 loader = DataLoader(VOC12(args.datadir, input_transform, target_transform), num_workers=args.num_workers, batch_size=args.batch_size, shuffle=True) if args.cuda: criterion = CrossEntropyLoss2d(weight.cuda()) else: criterion = CrossEntropyLoss2d(weight) #optimizer = Adam(model.parameters()) optimizer=SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3) ''' if args.model.startswith('FCN'): optimizer = SGD(model.parameters(), 1e-4, .9, 2e-5) if args.model.startswith('PSP'): optimizer = SGD(model.parameters(), 1e-2, .9, 1e-4) if args.model.startswith('Seg'): optimizer = SGD(model.parameters(), 1e-3, .9) ''' ''' if args.steps_plot > 0: board = Dashboard(args.port) ''' print(len(loader)) for epoch in range(1, args.num_epochs+1): epoch_loss = [] iteration=1 for step, (images, labels) in enumerate(loader): print("Iter:"+str(iteration)) iteration=iteration+1 if args.cuda: images = images.cuda() labels = labels.cuda() inputs = Variable(images) targets = Variable(labels) #change outputs = model(inputs) optimizer.zero_grad() loss = criterion(outputs, targets[:, 0]) loss.backward() optimizer.step() epoch_loss.append(loss.data[0]) if args.steps_plot > 0 and step % args.steps_plot == 0: image = inputs[0].cpu().data image[0] = image[0] * .229 + .485 image[1] = image[1] * .224 + .456 image[2] = image[2] * .225 + .406 ''' board.image(image, f'input (epoch: {epoch}, step: {step})') board.image(color_transform(outputs[0].cpu().max(0)[1].data), f'output (epoch: {epoch}, step: {step})') board.image(color_transform(targets[0].cpu().data), f'target (epoch: {epoch}, step: {step})') ''' #if args.steps_loss > 0 and step % args.steps_loss == 0: average = sum(epoch_loss) / len(epoch_loss) print("loss: "+str(average)+" epoch: "+str(epoch)+", step: "+str(step)) '''
seg_loader = torch.utils.data.DataLoader(MyBoxPixData(seg_train_dir, transform=True, crop=True, hflip=True, vflip=False, source=opt.q), batch_size=bsize, shuffle=True, num_workers=4, pin_memory=True) criterion_cls = nn.CrossEntropyLoss(weight=torch.FloatTensor(cls_label_weight)) criterion_cls.cuda() criterion_seg = CrossEntropyLoss2d(weight=torch.FloatTensor(seg_label_weight)) criterion_seg.cuda() optimizer_classifier = torch.optim.Adam(classifier.parameters(), lr=1e-3) optimizer_deconv = torch.optim.Adam(deconv.parameters(), lr=1e-3) optimizer_feature = torch.optim.Adam(feature.parameters(), lr=1e-4) segIter = iter(seg_loader) ibs = 0 for it in range(resume_ep + 1, iter_num): for ib, (data, lbl) in enumerate(cls_loader): # train with classification data inputs = Variable(data.float()).cuda() lbl = Variable(lbl.long()).cuda() feats = feature(inputs) output = classifier(feats)
batch_size=1, shuffle=True, num_workers=4, pin_memory=True) else: outputRoot = opt.out loader = torch.utils.data.DataLoader( MyTestData(dataRoot, transform=True), batch_size=1, shuffle=True, num_workers=4, pin_memory=True) """nets""" model = RFCN() if opt.param is None: vgg16 = torchvision.models.vgg16(pretrained=True) model.copy_params_from_vgg16(vgg16, copy_fc8=False, init_upscore=True) else: model.load_state_dict(torch.load(opt.param)) criterion = CrossEntropyLoss2d() optimizer = torch.optim.Adam(model.parameters(), 0.0001, betas=(0.5, 0.999)) model = model.cuda() mr_sal = MR.MR_saliency() if opt.phase == 'train': """train""" for it in range(iterNum): epoch_loss = [] for ib, data in enumerate(loader): # prior map _img = tensor2image(data[0][0]) pmap = mr_sal.saliency(_img).astype(float) / 255.0 pmap = 1.0 - pmap pmap = torch.unsqueeze(torch.FloatTensor(pmap), 0)
def train(args, model): # set model to training mode model.train() # prepare criterion weight = torch.ones(22) weight[0] = 0 train_loader = DataLoader(VOC12(root=args.datadir, train=True, input_transform=input_transform, target_transform=target_transform), num_workers=args.num_workers, batch_size=args.batch_size, shuffle=False) val_loader = DataLoader(VOC12(root=args.datadir, train=False, input_transform=input_transform, target_transform=target_transform), num_workers=args.num_workers, batch_size=args.batch_size, shuffle=False) # use Adam optimizer optimizer = Adam(model.parameters()) # for loss calculation, we still use CrossEntropyLoss if args.cuda: criterion = CrossEntropyLoss2d(weight.cuda()) else: criterion = CrossEntropyLoss2d(weight) # start training - epoch values start from 1 to make numbers look 'pretty' for epoch in range(1, args.num_epochs + 1): epoch_loss = [] print(len(train_loader)) for step, (images, labels) in enumerate(train_loader): if args.cuda: images = images.cuda() labels = labels.cuda() inputs = Variable(images) targets = Variable(labels) outputs = model(inputs) # refresh gradient before backprop optimizer.zero_grad() loss = criterion(outputs, targets[:, 0]) loss.backward() optimizer.step() epoch_loss.append(loss.item()) if (args.steps_loss > 0) and (step % args.steps_loss == 0): # print loss average = sum(epoch_loss) / len(epoch_loss) print(f'loss: {average} (epoch {epoch}, step {step})') # check for early stop if metrics.early_stop(use_cuda=args.cuda, output=outputs, target=targets[:, 0], metric="iou", threshold=0.8): # if we can stop early, save model and exit print( "Early stopping score exceeded threshold... saving model and ending training stage" ) if args.attention: filename = f'models/AttentionUNet-EarlyStop-{epoch:03}-{step:04}.pth' else: filename = f'models/UNet-EarlyStop-{epoch:03}-{step:04}.pth' torch.save(model.state_dict(), filename) return if (args.steps_save > 0) and (step % args.steps_save == 0): if args.attention: filename = f'models/AttentionUNet-{epoch:03}-{step:04}.pth' else: filename = f'models/UNet-{epoch:03}-{step:04}.pth' torch.save(model.state_dict(), filename) print(f'save: {filename} (epoch: {epoch}, step: {step})') ''' printing out pictures doesn't work on VScode :( _, outputs = torch.max(outputs, dim=1) #print(outputs.shape) outputs = outputs.unsqueeze(1) #print(outputs.shape) # show images plt.subplot(211) plt.imshow(make_grid(images.cpu()).permute(1,2,0).numpy()) plt.axis('off') plt.title('Images') # print labels plt.subplot(212) outputs = outputs.cpu().numpy()[:, :, :, :, np.newaxis] color_Label = np.dot(outputs == 0, cmap[0]) for i in range(1, cmap.shape[0]): color_Label += np.dot(outputs == i, cmap[i]) color_Label = color_Label.swapaxes(1,4) plt.imshow((make_grid(torch.tensor(color_Label.squeeze())).permute(1,2,0).numpy()).astype('uint8')) plt.axis('off') plt.title('Label') ''' # every epoch, check validation data's accuracy v_list = [] for _, (v_images, v_labels) in enumerate(val_loader): if args.cuda: v_images = v_images.cuda() v_labels = v_labels.cuda() v_inputs = Variable(v_images) v_targets = Variable(v_labels) v_outputs = model(v_inputs) # using cross-entropy loss # v_loss = criterion(v_outputs, v_targets[:, 0]) # v_list.append(v_loss.item()) # for IoU metric _, v_l = torch.max(v_outputs, dim=1) v_score = metrics.get_iou(labels=v_l, target=v_targets[:, 0]) v_list.append(v_score) v_average = sum(v_list) / len(v_list) print(f'validation average IoU: {v_average} (epoch {epoch})')
if __name__ == '__main__': models_path = os.path.join('./checkpoints', args.backend) #保存网络参数的路径 os.makedirs(models_path, exist_ok=True) train_loader, val_loader = get_dataloader() net, starting_epoch = build_network(args.snapshot, args.backend) optimizer = optim.Adam(net.parameters(), lr=args.start_lr) #优化器 # optimizer = optim.Adam(net.parameters()) # optimizer = optim.SGD(net.parameters(), lr=1e-3, momentum=0.9, # weight_decay=1e-4) scheduler = MultiStepLR( optimizer, milestones=[int(x) for x in args.milestones.split(',')]) seg_criterion = CrossEntropyLoss2d(weight=None) #损失函数 epoch_losses = [] net.train() weight_save_path = "checkpoints/densenet/PSPNet_last" #网络参数地址 #加载网络参数 try: net.load_state_dict(torch.load(weight_save_path)) print("加载成功") except: print('加载失败') #开始训练 Loss_list = [] Accuracy_list = [] for epoch in range(1 + starting_epoch, 1 + starting_epoch + args.epochs):
def train(args, model): model.train() weight = torch.ones(NUM_CLASSES) weight[0] = 0 loader = DataLoader(VOC12(args.datadir, input_transform, target_transform), num_workers=args.num_workers, batch_size=args.batch_size, shuffle=True) if args.cuda: criterion = CrossEntropyLoss2d(weight.cuda()) #criterion=torch.nn.BCEWithLogitsLoss() else: criterion = CrossEntropyLoss2d(weight) if args.model.startswith('FCN'): optimizer = SGD(model.parameters(), 1e-4, .9, 2e-5) if args.model.startswith('PSP'): optimizer = SGD(filter(lambda p: p.requires_grad, model.parameters()), 1e-2, 0.9, 1e-4) #optimizer = SGD(model.parameters(), 1e-2, .9, 1e-4) if args.model.startswith('Seg'): optimizer = SGD(filter(lambda p: p.requires_grad, model.parameters()), 1e-3, .9) print("Total images:", len(loader)) best_loss = 100 f = open("loss.txt", "a") for epoch in range(1, args.num_epochs + 1): epoch_loss = [] iteration = 1 for step, (images, labels) in enumerate(loader): print("Iter:" + str(iteration)) iteration = iteration + 1 if args.cuda: images = images.cuda() labels = labels.cuda() inputs = Variable(images) targets = Variable(labels) outputs = model(inputs) optimizer.zero_grad() loss = criterion(outputs, targets[:, 0]) loss.backward() optimizer.step() print(loss.data[0]) epoch_loss.append(loss.data[0]) if args.steps_loss > 0 and step > 0 and step % args.steps_loss == 0: average = sum(epoch_loss) / len(epoch_loss) epoch_loss = [] if best_loss > average: best_loss = average torch.save(model.state_dict(), "model_pspnet_VOC_2012_analysis.pth") print("Model saved!") f.write("loss: " + str(average) + " epoch: " + str(epoch) + ", step: " + str(step) + "\n") f.write("best loss: " + str(best_loss) + " epoch: " + str(epoch) + ", step: " + str(step) + "\n") print("loss: " + str(average) + " epoch: " + str(epoch) + ", step: " + str(step)) print("best loss: " + str(best_loss) + " epoch: " + str(epoch) + ", step: " + str(step)) print("Best loss: " + str(best_loss))