def train(args): # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset, config_file=args.config_file) loader = data_loader(data_path, is_transform=True, img_size=(args.img_rows, args.img_cols)) n_classes = loader.n_classes # must use 1 worker for AWS sagemaker without ipc="host" or larger shared memory size trainloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=1, shuffle=True) # Setup Model model = get_model(args.arch, n_classes) # Setup log dir / logging if not os.path.exists(args.log_dir): os.makedirs(args.log_dir) configure(args.log_dir) model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count())) model.cuda() optimizer = torch.optim.SGD(model.parameters(), lr=args.l_rate, momentum=0.9, weight_decay=5e-4) step = 0 for epoch in range(args.n_epoch): start_time = time.time() for i, (images, labels) in enumerate(trainloader): images = Variable(images.cuda()) labels = Variable(labels.cuda()) optimizer.zero_grad() outputs = model(images) loss = cross_entropy2d(outputs, labels) loss.backward() optimizer.step() log_value('Loss', loss.data[0], step) step += 1 if (i + 1) % 20 == 0: print("Epoch [%d/%d] Loss: %.4f" % (epoch + 1, args.n_epoch, loss.data[0]), flush=True) end_time = time.time() print('Epoch run time: %s' % (end_time - start_time)) torch.save( model, args.log_dir + "{}_{}_{}_{}.pt".format( args.arch, args.dataset, args.feature_scale, epoch))
def __init__(self, num_classes, pretrained=True, ignore_index=-1, weight=None, output_stride='16'): super(d_sunet7128, self).__init__() self.num_classes = num_classes sunet = stackedunet7128(output_stride=output_stride) sunet = torch.nn.DataParallel(sunet, device_ids=range( torch.cuda.device_count())).cuda() if pretrained: checkpoint = torch.load(sunet7128_path) sunet.load_state_dict(checkpoint['state_dict']) self.features = sunet.module.features for n, m in self.features.named_modules(): if 'bn' in n: m.momentum = mom_bn self.final = nn.Sequential(nn.Conv2d(2304, num_classes, kernel_size=1)) self.mceloss = cross_entropy2d(ignore=ignore_index, size_average=False, weight=weight)
def __init__(self, num_classes, pretrained=True, use_aux=True, ignore_index=-1, output_stride='16'): super(d_resnet18, self).__init__() self.use_aux = use_aux self.num_classes = num_classes resnet = models.resnet18() if pretrained: resnet.load_state_dict(torch.load(res18_path)) self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool) self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4 d = dilation[output_stride] if d > 1: for n, m in self.layer3.named_modules(): if '0.conv1' in n: m.dilation, m.padding, m.stride = (1, 1), (1, 1), (1, 1) elif 'conv1' in n: m.dilation, m.padding, m.stride = (d, d), (d, d), (1, 1) elif 'conv2' in n: m.dilation, m.padding, m.stride = (d, d), (d, d), (1, 1) elif 'downsample.0' in n: m.stride = (1, 1) for n, m in self.layer4.named_modules(): if '0.conv1' in n: m.dilation, m.padding, m.stride = (d, d), (d, d), (1, 1) elif 'conv1' in n: m.dilation, m.padding, m.stride = (2 * d, 2 * d), (2 * d, 2 * d), (1, 1) elif 'conv2' in n: m.dilation, m.padding, m.stride = (2 * d, 2 * d), (2 * d, 2 * d), (1, 1) elif 'downsample.0' in n: m.stride = (1, 1) for n, m in chain(self.layer0.named_modules(), self.layer1.named_modules(), self.layer2.named_modules(), self.layer3.named_modules(), self.layer4.named_modules()): if 'downsample.1' in n: m.momentum = mom_bn elif 'bn' in n: m.momentum = mom_bn self.final = nn.Sequential( nn.Conv2d(512, 512, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(512, momentum=mom_bn), nn.ReLU(inplace=True), nn.Dropout(0.1), nn.Conv2d(512, num_classes, kernel_size=1)) self.mceloss = cross_entropy2d(ignore=ignore_index, size_average=False)
def train(args): # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, is_transform=True, img_size=(args.img_rows, args.img_cols)) n_classes = loader.n_classes trainloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=4, shuffle=True) # Setup visdom for visualization vis = visdom.Visdom() loss_window = vis.line(X=torch.zeros((1, )).cpu(), Y=torch.zeros((1)).cpu(), opts=dict(xlabel='minibatches', ylabel='Loss', title='Training Loss', legend=['Loss'])) # Setup Model model = get_model(args.arch, n_classes) if torch.cuda.is_available(): model.cuda(0) test_image, test_segmap = loader[0] test_image = Variable(test_image.unsqueeze(0).cuda(0)) else: test_image, test_segmap = loader[0] test_image = Variable(test_image.unsqueeze(0)) optimizer = torch.optim.SGD(model.parameters(), lr=args.l_rate, momentum=0.99, weight_decay=5e-4) for epoch in range(args.n_epoch): for i, (images, labels) in enumerate(trainloader): if torch.cuda.is_available(): images = Variable(images.cuda(0)) labels = Variable(labels.cuda(0)) else: images = Variable(images) labels = Variable(labels) # iter = len(trainloader) * epoch + i # poly_lr_scheduler(optimizer, args.l_rate, iter) optimizer.zero_grad() outputs = model(images) loss = cross_entropy2d(outputs, labels) loss.backward() optimizer.step() vis.line(X=torch.ones((1, 1)).cpu() * i, Y=torch.Tensor([loss.data[0]]).unsqueeze(0).cpu(), win=loss_window, update='append') if (i + 1) % 20 == 0: print("Epoch [%d/%d] Loss: %.4f" % (epoch + 1, args.n_epoch, loss.data[0])) # test_output = model(test_image) # predicted = loader.decode_segmap(test_output[0].cpu().data.numpy().argmax(0)) # target = loader.decode_segmap(test_segmap.numpy()) # vis.image(test_image[0].cpu().data.numpy(), opts=dict(title='Input' + str(epoch))) # vis.image(np.transpose(target, [2,0,1]), opts=dict(title='GT' + str(epoch))) # vis.image(np.transpose(predicted, [2,0,1]), opts=dict(title='Predicted' + str(epoch))) torch.save( model, "{}_{}_{}_{}.pkl".format(args.arch, args.dataset, args.feature_scale, epoch))
def train(args, guotai_config): # Setup Dataloader print('###### Step One: Setup Dataloader') data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) # For 2D dataset keep is_transform True # loader = data_loader(data_path, is_transform=True, img_size=(args.img_rows, args.img_cols)) # For 3D dataset keep is_transform False # loader = data_loader(data_path, is_transform=False, img_size=(args.img_rows, args.img_cols)) if args.dataset == 'brats17_loader_guotai': config_data = config['data'] # print(config_data) # config_net = config['network'] config_train = config['training'] random.seed(config_train.get('random_seed', 1)) assert (config_data['with_ground_truth']) # net_type = config_net['net_type'] # net_name = config_net['net_name'] # class_num = config_net['class_num'] # batch_size = config_data.get('batch_size', 5) dataloader_guotai = DataLoader(config_data) dataloader_guotai.load_data() loader = data_loader(dataloader_guotai) elif args.dataset == 'brats17_loader': loader = data_loader(data_path, is_transform=False, img_size=(args.img_rows, args.img_cols)) else: loader = data_loader(data_path, is_transform=True, img_size=(args.img_rows, args.img_cols)) n_classes = args.n_classes trainloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=4, shuffle=True) # Setup Model print('###### Step Two: Setup Model') model = get_model(args.arch, n_classes) if args.pretrained_path != 'empty': model = torch.load(args.pretrained_path) #model = torch.load('/home/donghao/Desktop/donghao/isbi2019/code/fast_segmentation_code/runs/bisenet3Dbrain_brats17_loader_1_251_3020_min.pkl') #model = torch.load('/home/donghao/Desktop/donghao/isbi2019/code/fast_segmentation_code/runs/2177/bisenet3Dbrain_brats17_loader_1_293_min.pkl') #model = torch.load('/home/donghao/Desktop/donghao/isbi2019/code/fast_segmentation_code/runs/9863/FCDenseNet57_brats17_loader_1_599.pkl') # model = if torch.cuda.is_available(): model.cuda(0) test_image, test_segmap = loader[0] test_image = Variable(test_image.unsqueeze(0).cuda(0)) else: test_image, test_segmap = loader[0] test_image = Variable(test_image.unsqueeze(0)) log('The optimizer is Adam') log('The learning rate is {}'.format(args.l_rate)) # optimizer = torch.optim.Adam(model.parameters(), lr=1e-1) # optimizer = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=0.99) optimizer = torch.optim.Adam(model.parameters(), lr=args.l_rate) # optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) # Train Model print('###### Step Three: Training Model') epoch_loss_array_total = np.zeros([1, 2]) for epoch in range(args.n_epoch): img_counter = 1 loss_sum = 0 for i, (images, labels) in enumerate(trainloader): img_counter = img_counter + 1 if torch.cuda.is_available(): images = Variable(images.cuda(0)) labels = Variable(labels.cuda(0)) else: images = Variable(images) labels = Variable(labels) optimizer.zero_grad() # log('The maximum value of input image is {}'.format(images.max())) # print(images) outputs = model(images) # print(outputs) if args.arch == 'bisenet3Dbrain' or args.arch == 'unet3d_cls' or args.arch == 'FCDenseNet57' or args.arch == 'FCDenseNet103': loss = cross_entropy3d(outputs, labels) elif args.arch == 'unet3d_res': labels = labels * 40 labels = labels + 1 log('The unique value of labels are {}'.format( np.unique(labels))) log('The maximum of outputs are {}'.format(outputs.max())) log('The size of output is {}'.format(outputs.size())) log('The size of labels is {}'.format(labels.size())) loss = nn.L1Loss() labels = labels.type(torch.cuda.FloatTensor) outputs = torch.squeeze(outputs, dim=1) loss = loss(outputs, labels) else: loss = cross_entropy2d(outputs, labels) loss.backward() optimizer.step() loss_sum = loss_sum + torch.Tensor([loss.data]).unsqueeze(0).cpu() avg_loss = loss_sum / img_counter avg_loss_array = np.array(avg_loss) epoch_loss_array_total = np.concatenate( (epoch_loss_array_total, [[avg_loss_array[0][0], epoch]]), axis=0) print('The current loss of epoch', epoch, 'is', avg_loss_array[0][0]) # training model will be saved log('The variable avg_loss_array is {}'.format(avg_loss_array)) writer.add_scalar('train_main_loss', avg_loss_array[0][0], epoch) if epoch % 10 == 0: torch.save( model, "runs/{}/{}_{}_{}_{}.pkl".format(rand_int, args.arch, args.dataset, args.feature_scale, epoch)) # I guess the shape is (epoch, 2) log('epoch_loss_array_total is {}'.format(epoch_loss_array_total)) # The shape of epoch_loss_array_total is (epoch, 2) log('the shape of epoch_loss_array_total is {}'.format( epoch_loss_array_total.shape)) epoch_loss_array_total = np.delete(arr=epoch_loss_array_total, obj=0, axis=0) log('the shape of epoch_loss_array_total after removal is {}'.format( epoch_loss_array_total.shape)) loss_min_indice = np.argmin(epoch_loss_array_total, axis=0) log('The loss_min_indice is {}'.format(loss_min_indice)) torch.save( model, "runs/{}/{}_{}_{}_{}_min.pkl".format(rand_int, args.arch, args.dataset, args.feature_scale, loss_min_indice[0])) sys.stdout = orig_stdout f.close()
def train(args): # Setup Dataloader data_loader = get_loader("mangrove") data_path = get_data_path("mangrove") loader = data_loader(data_path, img_size=args.img_size) n_classes = loader.n_classes n_channels = loader.n_channels trainloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=4, shuffle=True) # Setup visdom for visualization vis = visdom.Visdom() loss_window = vis.line(X=torch.zeros((1, )).cpu(), Y=torch.zeros((1)).cpu(), opts=dict(xlabel='minibatches', ylabel='Loss', title='Training Loss', legend=['Loss'])) # Setup Model # default: coralnet model = get_model("coralnet", n_classes, in_channels=n_channels) if torch.cuda.is_available(): model.cuda(0) test_image, test_segmap = loader[0] test_image = Variable(test_image.unsqueeze(0).cuda(0)) else: print("CUDA Error.") test_image, test_segmap = loader[0] test_image = Variable(test_image.unsqueeze(0)) # optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate, momentum=0.99, weight_decay=5e-4) optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate) for epoch in range(args.n_epoch + 1): for i, (images, labels) in enumerate(trainloader): if torch.cuda.is_available(): images = Variable(images.cuda(0)) labels = Variable(labels.cuda(0)) else: images = Variable(images) labels = Variable(labels) iter = len(trainloader) * epoch + i # poly_lr_scheduler(optimizer, args.learning_rate, iter) optimizer.zero_grad() outputs = model(images) loss = cross_entropy2d(outputs, labels) loss.backward() optimizer.step() vis.line(X=torch.ones((1, 1)).cpu() * i, Y=torch.Tensor([loss.data[0]]).unsqueeze(0).cpu(), win=loss_window, update='append') if (i + 1) % 20 == 0: print("Epoch [%d/%d] Loss: %.4f" % (epoch + 1, args.n_epoch, loss.data[0])) if epoch % 50 == 0: torch.save( model, "training/{}_{}_{}_{}.pkl".format("mangrove_multi", "linknet", args.feature_scale, epoch + 1))
def train(args): # writer = SummaryWriter('exp') # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, is_transform=True, img_size=(args.img_rows, args.img_cols)) n_classes = loader.n_classes trainloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=4, shuffle=True) # Setup visdom for visualization # vis = visdom.Visdom() # loss_window = vis.line(X=torch.zeros((1,)).cpu(), # Y=torch.zeros((1)).cpu(), # opts=dict(xlabel='epoch', # ylabel='Loss', # title='Training Loss', # legend=['Loss'])) # Setup Model model = get_model(args.arch, n_classes) if torch.cuda.is_available(): model.cuda(0) test_image, test_segmap = loader[0] test_image = Variable(test_image.unsqueeze(0).cuda(0)) else: test_image, test_segmap = loader[0] test_image = Variable(test_image.unsqueeze(0)) # optimizer = torch.optim.SGD(model.parameters(), lr=args.l_rate, momentum=0.99, weight_decay=5e-4) optimizer = torch.optim.Adam(model.parameters(), lr=args.l_rate * 100) for epoch in range(args.n_epoch): img_counter = 1 loss_sum = 0 for i, (images, labels) in enumerate(trainloader): img_counter = img_counter + 1 # print(img_counter) # print('i value: ', i) if torch.cuda.is_available(): images = Variable(images.cuda(0)) labels = Variable(labels.cuda(0)) else: images = Variable(images) labels = Variable(labels) # iter = len(trainloader)*epoch + i # poly_lr_scheduler(optimizer, args.l_rate, iter) optimizer.zero_grad() outputs = model(images) loss = cross_entropy2d(outputs, labels) loss.backward() optimizer.step() # vis.line( # X=torch.ones((1, 1)).cpu() * i, # Y=torch.Tensor([loss.data[0]]).unsqueeze(0).cpu(), # win=loss_window, # update='append') if (i + 1) % 20 == 0: print("Epoch [%d/%d] Loss: %.4f" % (epoch + 1, args.n_epoch, loss.data[0])) # print('The number of img_counter is ', img_counter) loss_sum = loss_sum + torch.Tensor([loss.data[0] ]).unsqueeze(0).cpu() avg_loss = loss_sum / img_counter avg_loss_array = np.array(avg_loss) # vis.line( # X=torch.ones((1, 1)).cpu() * epoch, # Y=avg_loss, # win=loss_window, # update='append') # writer.add_scalar('train_main_loss', avg_loss, epoch) test_output = model(test_image) predicted = loader.decode_segmap( test_output[0].cpu().data.numpy().argmax(0)) target = loader.decode_segmap(test_segmap.numpy()) # if epoch == 1: # vis.image(test_image[0].cpu().data.numpy(), opts=dict(title='Test Epoch' + str(epoch))) # vis.image(np.transpose(target, [2,0,1]), opts=dict(title='GT Epoch' + str(epoch))) # vis.image(np.transpose(predicted, [2,0,1]), opts=dict(title='Predicted Epoch' + str(epoch))) torch.save( model, "{}_{}_{}_{}.pkl".format(args.arch, args.dataset, args.feature_scale, epoch))
def train(args): # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, is_transform=True, img_size=(args.img_rows, args.img_cols)) n_classes = loader.n_classes trainloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=4, shuffle=True) # Setup Model net_features, net_segmenter = get_model(name=args.arch, n_classes=n_classes) netD = _netD() netG = _netG(nz=nz) padder = padder_layer(pad_size=100) criterion_gan = nn.BCELoss() input = torch.FloatTensor(args.batch_size_gan, 3, args.img_rows, args.img_cols) noise = torch.FloatTensor(args.batch_size_gan, nz, 1, 1) fixed_noise = torch.FloatTensor(args.batch_size_gan, nz, 1, 1).normal_(0, 1) label = torch.FloatTensor(args.batch_size_gan) real_label = .9 fake_label = .1 G = args.gpu if torch.cuda.is_available(): input, label = input.cuda(G), label.cuda(G) noise, fixed_noise = noise.cuda(G), fixed_noise.cuda(G) fixed_noise = Variable(fixed_noise) if torch.cuda.is_available(): net_features.cuda(G) net_segmenter.cuda(G) padder.cuda(G) netD.cuda(G) netG.cuda(G) criterion_gan.cuda(G) optimizerS = torch.optim.SGD(list(net_features.parameters()) + list(net_segmenter.parameters()), lr=args.l_rate, momentum=0.99, weight_decay=5e-4) optimizerD = torch.optim.Adam(list(net_features.parameters()) + list(netD.parameters()), lr=args.l_rate, betas=(args.beta1, 0.999)) optimizerG = torch.optim.Adam(netG.parameters(), lr=args.l_rate, betas=(args.beta1, 0.999)) SS = 1 GAN = 0 for epoch in range(args.n_epoch): for i, (images, labels) in enumerate(trainloader): real_cpu, label_cpu = data batch_size = real_cpu.size(0) if torch.cuda.is_available(): real_cpu = real_cpu.cuda(G) label_cpu = label_cpu.cuda(G) input.resize_as_(real_cpu).copy_(real_cpu) inputv = Variable(input) labelv_semantic = Variable(label_cpu) # iter = len(trainloader)*epoch + i # poly_lr_scheduler(optimizer, args.l_rate, iter) if GAN: ############################ # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z))) ########################### optimizerD.zero_grad() netD.zero_grad() net_features.zero_grad() # train with real label.resize_(batch_size).fill_(real_label) labelv = Variable(label) inputv_feats = net_features(inputv) output = netD(inputv_feats) errD_real = criterion_gan(output, labelv) errD_real.backward() D_x = output.data.mean() # train with fake noise.resize_(batch_size, nz, 1, 1).normal_(0, 1) noisev = Variable(noise) fake = netG(noisev) labelv = Variable(label.fill_(fake_label)) fake_feats = net_features(fake.detach()) output = netD(fake_feats) errD_fake = criterion_gan(output, labelv) errD_fake.backward() D_G_z1 = output.data.mean() errD = errD_real + errD_fake optimizerD.step() ############################ # (2) Update G network: maximize log(D(G(z))) ########################### optimizerG.zero_grad() netG.zero_grad() label.resize_(batch_size).fill_(real_label) labelv = Variable(label.fill_( real_label)) # fake labels are real for generator cost noise.resize_(batch_size, nz, 1, 1).normal_(0, 1) noise2v = Variable(noise) fake2 = netG(noise2v) fake_feats2 = net_features(fake2) output = netD(fake_feats2) errG = criterion_gan(output, labelv) errG.backward() D_G_z2 = output.data.mean() optimizerG.step() if SS: optimizerS.zero_grad() net_features.zero_grad() net_segmenter.zero_grad() padded_images = padder(inputv) features = net_features(padded_images) outputs = net_segmenter(features) loss = cross_entropy2d(outputs, labelv_semantic) loss.backward() optimizerS.step() if (i + 1) % 20 == 0: print("Iter [%d/%d], Epoch [%d/%d] Loss: %.4f" % (i + 1, len(trainloader), epoch + 1, args.n_epoch, loss.data[0])) # test_output = model(test_image) # predicted = loader.decode_segmap(test_output[0].cpu().data.numpy().argmax(0)) # target = loader.decode_segmap(test_segmap.numpy()) # vis.image(test_image[0].cpu().data.numpy(), opts=dict(title='Input' + str(epoch))) # vis.image(np.transpose(target, [2,0,1]), opts=dict(title='GT' + str(epoch))) # vis.image(np.transpose(predicted, [2,0,1]), opts=dict(title='Predicted' + str(epoch))) torch.save((net_features, net_segmenter), "{}_{}_{}_{}.pkl".format(args.arch, args.dataset, args.feature_scale, epoch))
def train(args): # Setup Dataloader if args.dataset in ['imagenet', 'folder', 'lfw']: # folder dataset assert (args.img_cols == args.img_rows) dataset = dset.ImageFolder(root=args.dataroot, transform=transforms.Compose([ transforms.Scale(args.img_cols), transforms.CenterCrop(args.img_cols), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ])) elif args.dataset in ['pascal', 'camvid']: data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) dataset = data_loader(data_path, is_transform=True, img_size=(args.img_rows, args.img_cols)) n_classes = dataset.n_classes trainloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, num_workers=4, shuffle=True) ############################ # Setup the Models net_features, net_segmenter = get_model(name=args.arch, n_classes=n_classes) netD = _netD() netG = _netG(nz=nz) padder = padder_layer(pad_size=100) ############################ ############################ ### Initialization: if args.netD != '': netD.load_state_dict(torch.load(os.path.join(args.outf, args.netD))) if args.netG != '': netG.load_state_dict(torch.load(os.path.join(args.outf, args.netG))) if args.net_features != '': netD.load_state_dict( torch.load(os.path.join(args.outf, args.net_features))) if args.net_segmenter != '': netD.load_state_dict( torch.load(os.path.join(args.outf, args.net_segmenter))) ############################ criterion_gan = nn.BCELoss() input = torch.FloatTensor(args.batch_size_gan, 3, args.img_rows, args.img_cols) noise = torch.FloatTensor(args.batch_size_gan, nz, 1, 1) fixed_noise = torch.FloatTensor(args.batch_size_gan, nz, 1, 1).normal_(0, 1) label = torch.FloatTensor(args.batch_size_gan) real_label = .9 fake_label = .1 G = args.gpu if torch.cuda.is_available(): input, label = input.cuda(G), label.cuda(G) noise, fixed_noise = noise.cuda(G), fixed_noise.cuda(G) fixed_noise = Variable(fixed_noise) if torch.cuda.is_available(): net_features.cuda(G) net_segmenter.cuda(G) padder.cuda(G) netD.cuda(G) netG.cuda(G) criterion_gan.cuda(G) optimizerS = torch.optim.SGD(list(net_features.parameters()) + list(net_segmenter.parameters()), lr=args.l_rate, momentum=0.99, weight_decay=5e-4) optimizerD = torch.optim.Adam(list(net_features.parameters()) + list(netD.parameters()), lr=args.l_rate, betas=(args.beta1, 0.999)) optimizerG = torch.optim.Adam(netG.parameters(), lr=args.l_rate, betas=(args.beta1, 0.999)) SS = 1 GAN = 0 for epoch in range(args.n_epoch): for i, data in enumerate(trainloader): real_cpu, label_cpu = data batch_size = real_cpu.size(0) if torch.cuda.is_available(): real_cpu = real_cpu.cuda(G) label_cpu = label_cpu.cuda(G) input.resize_as_(real_cpu).copy_(real_cpu) inputv = Variable(input) labelv_semantic = Variable(label_cpu) # iter = len(trainloader)*epoch + i # poly_lr_scheduler(optimizer, args.l_rate, iter) if GAN: ############################ # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z))) ########################### optimizerD.zero_grad() net_features.zero_grad() netD.zero_grad() # train with real label.resize_(batch_size).fill_(real_label) labelv = Variable(label) inputv_feats = net_features(inputv) output = netD(inputv_feats) errD_real = criterion_gan(output, labelv) errD_real.backward() D_x = output.data.mean() # train with fake noise.resize_(batch_size, nz, 1, 1).normal_(0, 1) noisev = Variable(noise) fake = netG(noisev) labelv = Variable(label.fill_(fake_label)) fake_feats = net_features(fake.detach()) output = netD(fake_feats) errD_fake = criterion_gan(output, labelv) errD_fake.backward() D_G_z1 = output.data.mean() errD = errD_real + errD_fake optimizerD.step() ############################ # (2) Update G network: maximize log(D(G(z))) ########################### optimizerG.zero_grad() netG.zero_grad() label.resize_(batch_size).fill_(real_label) labelv = Variable(label.fill_( real_label)) # fake labels are real for generator cost noise.resize_(batch_size, nz, 1, 1).normal_(0, 1) noise2v = Variable(noise) fake2 = netG(noise2v) fake_feats2 = net_features(fake2) output = netD(fake_feats2) errG = criterion_gan(output, labelv) errG.backward() D_G_z2 = output.data.mean() optimizerG.step() if SS: optimizerS.zero_grad() net_features.zero_grad() net_segmenter.zero_grad() padded_images = padder(inputv) features = net_features(padded_images) outputs = net_segmenter(features) loss = cross_entropy2d(outputs, labelv_semantic) loss.backward() optimizerS.step() ###################################################### ### Loss Report: ###################################################### if SS: if (i + 1) % 20 == 0: print("Iter [%d/%d], Epoch [%d/%d] Loss: %.4f" % (i + 1, len(trainloader), epoch + 1, args.n_epoch, loss.data[0])) if GAN: if (i) % 100 == 0 and i > 0: print( '[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f' % (epoch, args.n_epoch, i, len(trainloader), errD.data[0], errG.data[0], D_x, D_G_z1, D_G_z2)) if i % 500 == 0 and i > 0: vutils.save_image(real_cpu, '%s/real_samples.png' % args.outf, normalize=True) fake = netG(fixed_noise) vutils.save_image(fake.data, '%s/fake_samples_epoch_%03d_%03d.png' % (args.outf, epoch, i), normalize=True) if i % 3000 == 0 and i > 0: torch.save( net_features.state_dict(), '%s/net_features_epoch_%d_%d.pth' % (args.outf, epoch, i)) torch.save( netG.state_dict(), '%s/netG_epoch_%d_%d.pth' % (args.outf, epoch, i)) torch.save( netD.state_dict(), '%s/netD_epoch_%d_%d.pth' % (args.outf, epoch, i)) ###################################################### ### Do checkpointing: ###################################################### torch.save((net_features, net_segmenter), "{}_{}_{}_{}.pkl".format(args.arch, args.dataset, args.feature_scale, epoch)) continue torch.save(net_features.state_dict(), '%s/net_features_epoch_%d.pth' % (args.outf, epoch)) torch.save(netG.state_dict(), '%s/netG_epoch_%d.pth' % (args.outf, epoch)) torch.save(netD.state_dict(), '%s/netD_epoch_%d.pth' % (args.outf, epoch))
def train(model): if model == 'unet': model = unet(feature_scale=feature_scale, n_classes=n_classes, is_batchnorm=True, in_channels=3, is_deconv=True) if model == 'segnet': model = segnet(n_classes=n_classes, in_channels=3, is_unpooling=True) if model == 'fcn32': model = fcn32s(n_classes=n_classes) vgg16 = models.vgg16(pretrained=True) model.init_vgg16_params(vgg16) if model == 'fcn16': model = fcn16s(n_classes=n_classes) vgg16 = models.vgg16(pretrained=True) model.init_vgg16_params(vgg16) if model == 'fcn8': model = fcn8s(n_classes=n_classes) vgg16 = models.vgg16(pretrained=True) model.init_vgg16_params(vgg16) pascal = pascalVOCLoader(data_path, is_transform=True, img_size=img_rows) trainloader = data.DataLoader(pascal, batch_size=batch_size, num_workers=4) if torch.cuda.is_available(): model.cuda(0) optimizer = torch.optim.SGD(model.parameters(), lr=l_rate, momentum=0.99, weight_decay=5e-4) test_image, test_segmap = pascal[0] test_image = Variable(test_image.unsqueeze(0).cuda(0)) vis = visdom.Visdom() for epoch in range(n_epoch): for i, (images, labels) in enumerate(trainloader): if torch.cuda.is_available(): images = Variable(images.cuda(0)) labels = Variable(labels.cuda(0)) else: images = Variable(images) labels = Variable(labels) optimizer.zero_grad() outputs = model(images) loss = cross_entropy2d(outputs, labels) loss.backward() optimizer.step() if (i + 1) % 20 == 0: print("Epoch [%d/%d] Loss: %.4f" % (epoch + 1, n_epoch, loss.data[0])) test_output = model(test_image) predicted = pascal.decode_segmap( test_output[0].cpu().data.numpy().argmax(0)) target = pascal.decode_segmap(test_segmap.numpy()) vis.image(test_image[0].cpu().data.numpy(), opts=dict(title='Input' + str(epoch))) vis.image(np.transpose(target, [2, 0, 1]), opts=dict(title='GT' + str(epoch))) vis.image(np.transpose(predicted, [2, 0, 1]), opts=dict(title='Predicted' + str(epoch))) torch.save(model, "unet_voc_" + str(feature_scale) + ".pkl")
def train(args): # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, split=args.split, is_transform=True, img_size=(args.img_rows, args.img_cols)) n_classes = loader.n_classes trainloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=4, shuffle=True) # Setup Model pre_trained = True model = get_model(args.arch, n_classes, pre_trained=pre_trained) ############################################ # Random weights initialization if not pre_trained: model.apply(weights_init) ############################################ if args.restore_from != '': print('\n' + '-' * 40) model = torch.load(args.restore_from) print('Restored the trained network.') print('-' * 40) if torch.cuda.is_available(): model.cuda(args.gpu) optimizer = torch.optim.SGD(model.parameters(), lr=args.l_rate, momentum=0.99, weight_decay=5e-4) for epoch in range(args.n_epoch): for i, (images, labels) in enumerate(trainloader): if torch.cuda.is_available(): images = Variable(images.cuda(args.gpu)) labels = Variable(labels.cuda(args.gpu)) else: images = Variable(images) labels = Variable(labels) #iter = len(trainloader)*epoch + i adjust_learning_rate(optimizer, args.l_rate, epoch) optimizer.zero_grad() outputs = model(images) loss = cross_entropy2d(outputs, labels) loss.backward() optimizer.step() if (i+1) % 20 == 0: print("Iter [%d/%d], Epoch [%d/%d] Loss: %.4f" % (i+1, len(trainloader), epoch+1, args.n_epoch, loss.data[0])) # test_output = model(test_image) # predicted = loader.decode_segmap(test_output[0].cpu().data.numpy().argmax(0)) # target = loader.decode_segmap(test_segmap.numpy()) # vis.image(test_image[0].cpu().data.numpy(), opts=dict(title='Input' + str(epoch))) # vis.image(np.transpose(target, [2,0,1]), opts=dict(title='GT' + str(epoch))) # vis.image(np.transpose(predicted, [2,0,1]), opts=dict(title='Predicted' + str(epoch))) if args.restore_from != '': torch.save(model, "./{}/{}_{}_{}_from_{}.pkl".format(args.save_folder, args.arch, args.dataset, epoch, args.restore_from)) else: torch.save(model, "./{}/{}_{}_{}.pkl".format(args.save_folder, args.arch, args.dataset, epoch))
def train(args): ################### pre_trained = args.pre_trained # default='gt' ################### ############################################ if pre_trained == 'gt': # When pre_trained = 'gt', i.e. when using supervised image-net weights, images were normalized with image-net mean. image_transform = None else: # When pre_trained = 'self' or 'no', i.e. in the self-supervised case, or unsupervised case, the input images are normalized this way: image_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, split=args.split, is_transform=True, img_size=(args.img_rows, args.img_cols), image_transform=image_transform) n_classes = loader.n_classes trainloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=4, shuffle=True) ############################################ ############################################ # Setup Model if args.netF_path == '' or args.netS_path == '': if pre_trained == 'gt': netF, netS = get_model(args.arch, n_classes, pre_trained=True) print('*' * 60) print('*' * 60) print('*' * 60) print('Training using a GT-supervised pre-trained model ...') print('*' * 60) print('*' * 60) print('*' * 60 + '\n') elif pre_trained == 'self': netF, netS = get_model(args.arch, n_classes, pre_trained=False) assert (args.model_path != '') # Loading a self-supervised trained model (ss_model) print('x' * 60) print('x' * 60) print('x' * 60) print('Training a using a self-supervised pre-trained model ...') ss_model = torch.load(args.model_path) netF.init_alex_params(ss_model) #netS.init_alex_params(ss_model) netS.apply(weights_init) print('Restored the self-supervised pre_trained model {}'.format(args.model_path)) print('x' * 60) print('x' * 60) print('x' * 60 + '\n') else: # pre_trained == 'no': netF, netS = get_model(args.arch, n_classes, pre_trained=False) # Random weight initialization netF.apply(weights_init) netS.apply(weights_init) print('o' * 60) print('o' * 60) print('o' * 60) print('Training without any pre-trained model ...') print('o' * 60) print('o' * 60) print('o' * 60 + '\n') padder = padder_layer(pad_size=100) ############################################ ############################################ # If resuming the training from a saved model if args.netF_path != '': print('\n' + '-' * 40) netF = torch.load(args.netF_path) print('Resuming the training netF from model in {}'.format(args.netF_path)) if args.netS_path != '': netS = torch.load(args.netS_path) print('Resuming the training netS from model in {}'.format(args.netS_path)) print('-' * 40) ############################################ ############################################ # Porting the networks to CUDA if torch.cuda.is_available(): netF.cuda(args.gpu) netS.cuda(args.gpu) padder.cuda(args.gpu) ############################################ ############################################ # Defining the optimizer over the network parameters optimizerSS = torch.optim.SGD([{'params': netF.features.parameters()}, {'params': netF.classifier.parameters(), 'lr':10*args.l_rate}, {'params': netS.parameters(), 'lr':20*args.l_rate}], lr=args.l_rate, momentum=0.99, weight_decay=5e-4) optimizerSS_init = copy.deepcopy(optimizerSS) ############################################ ############################################ # TRAINING: for epoch in range(args.n_epoch): for i, (images, labels) in enumerate(trainloader): ###################### # Porting the data to Autograd variables and CUDA (if available) if torch.cuda.is_available(): images = Variable(images.cuda(args.gpu)) labels = Variable(labels.cuda(args.gpu)) else: images = Variable(images) labels = Variable(labels) ###################### # Scheduling the learning rate if args.pre_trained=='no': adjust_learning_rate(optimizerSS, args.l_rate, epoch, step=20) else: adjust_learning_rate_v2(optimizerSS, optimizerSS_init, epoch, step=20) ###################### # Setting the gradients to zero at each iteration optimizerSS.zero_grad() netF.zero_grad() netS.zero_grad() ###################### # Passing the data through the networks padded_images = padder(images) feature_maps = netF(padded_images) score_maps = netS(feature_maps) outputs = F.upsample(score_maps, labels.size()[1:], mode='bilinear') ###################### # Computing the loss and doing back-propagation loss = cross_entropy2d(outputs, labels) loss.backward() ###################### # Updating the parameters optimizerSS.step() if (i+1) % 20 == 0: print("Iter [%d/%d], Epoch [%d/%d] Loss: %.4f" % (i+1, len(trainloader), epoch+1, args.n_epoch, loss.data[0])) # test_output = model(test_image) # predicted = loader.decode_segmap(test_output[0].cpu().data.numpy().argmax(0)) # target = loader.decode_segmap(test_segmap.numpy()) # vis.image(test_image[0].cpu().data.numpy(), opts=dict(title='Input' + str(epoch))) # vis.image(np.transpose(target, [2,0,1]), opts=dict(title='GT' + str(epoch))) # vis.image(np.transpose(predicted, [2,0,1]), opts=dict(title='Predicted' + str(epoch))) torch.save(netF, "./{}/netF_{}_{}_{}.pkl".format(args.save_folder, args.arch, args.dataset, epoch)) torch.save(netS, "./{}/netS_{}_{}_{}.pkl".format(args.save_folder, args.arch, args.dataset, epoch))
def train(args): # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, is_transform=True, img_size=(args.img_rows, args.img_cols)) n_classes = loader.n_classes trainloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=4, shuffle=True) # Setup visdom for visualization if args.visdom: vis = visdom.Visdom() loss_window = vis.line(X=torch.zeros((1, )).cpu(), Y=torch.zeros((1)).cpu(), opts=dict(xlabel='minibatches', ylabel='Loss', title='Training Loss', legend=['Loss'])) # Setup Model model = get_model(args.arch, n_classes) model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count())) model.cuda() optimizer = torch.optim.SGD(model.parameters(), lr=args.l_rate, momentum=0.99, weight_decay=5e-4) for epoch in range(args.n_epoch): for i, (images, labels) in enumerate(trainloader): images = Variable(images.cuda()) labels = Variable(labels.cuda()) optimizer.zero_grad() outputs = model(images) loss = cross_entropy2d(outputs, labels) loss.backward() optimizer.step() if args.visdom: vis.line(X=torch.ones((1, 1)).cpu() * i, Y=torch.Tensor([loss.data[0]]).unsqueeze(0).cpu(), win=loss_window, update='append') if (i + 1) % 20 == 0: print("%s / Epoch [%d/%d] Loss: %.4f" % (args.arch, epoch + 1, args.n_epoch, loss.data[0])) # torch.save(model, "{}_{}_{}_{}.pkl".format(args.arch, args.dataset, args.feature_scale, epoch)) torch.save( model, "{}_{}_{}.pkl".format(args.arch, args.dataset, args.feature_scale))
def train(args, model, optimizer, dataset, episode=0): trainloader = data.DataLoader(dataset, batch_size=args.batch_size, num_workers=8, shuffle=True, drop_last=True) class_weight = Variable(dataset.class_weight.cuda()) lr = args.l_rate n_epoch = args.n_epoch optimizer.param_groups[0]['lr'] = args.l_rate model.train() # Setup visdom for visualization if args.visdom: vis = visdom.Visdom(port=args.visdom) loss_window = vis.line(X=np.column_stack((np.zeros((1,)))), Y=np.column_stack((np.zeros((1)))), opts=dict(xlabel='epoch', ylabel='Loss', title=args.mode + '_' + args.exp_name + '_Episode_' + str(episode), legend=['Train Loss'])) t1 = time.time() start_epoch = args.start_epoch if episode == args.start_episode else 0 best_iou = -100.0 save_interval = int(floor(n_epoch*args.save_percent)) for epoch in range(1 + start_epoch, n_epoch + 1): utils.adjust_learning_rate(optimizer, args.l_rate, args.lr_decay, epoch - 1, 1) for i, (images, labels, image_name) in enumerate(trainloader): images = Variable(images.cuda()) labels = Variable(labels.cuda(async=True)) optimizer.zero_grad() outputs = model(images) loss = cross_entropy2d(outputs, labels, class_weight) loss.backward() optimizer.step() if epoch % (save_interval*args.eval_interval) == 0: gts, preds, uncts = test(args, model=model, split='val') model.train() _, score = eval_metrics(args, gts, preds, verbose=False) print 'val Mean IoU: ', score['Mean IoU : \t'] if score['Mean IoU : \t'] >= best_iou: best_iou = score['Mean IoU : \t'] state = {'episode': episode, 'epoch': epoch, 'model_state': model.state_dict(), 'optimizer_state' : optimizer.state_dict(),} print "update best model {}".format(best_iou) torch.save(state, "checkpoint/{}/{}_{}_{}_best_model.pkl".format(\ args.exp_name, args.arch, 'camvid', episode)) utils.adjust_learning_rate(optimizer, args.l_rate, args.lr_decay, epoch - 1, 1) if epoch % save_interval == 0: print 'data_size : ', len(dataset) state = { 'episode' : episode, 'epoch': epoch, 'arch': args.arch, 'loss': loss.data[0], 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict() } torch.save(state, 'checkpoint/{}/{}_{}_{}.pth.tar'.format(\ args.exp_name, args.arch, episode, epoch)) print("Epoch [%d/%d] Loss: %.4f lr:%.4f" % (epoch, n_epoch, loss.data[0], optimizer.param_groups[0]['lr'] )) t2 = time.time() print save_interval, 'epoch time :', t2 - t1 t1 = time.time() if args.visdom: vis.line( X=np.column_stack((np.ones((1,)) * epoch)), Y=np.column_stack((np.array([loss.data[0]]))), win=loss_window, update='append') return model, optimizer