def main(): # parse options parser = TrainOptions() opts = parser.parse() # daita loader print('\n--- load dataset ---') os.makedirs(opts.dataroot, exist_ok=True) dataset = torchvision.datasets.CIFAR10(opts.dataroot, train=True, download=True, transform= transforms.Compose([ transforms.Resize(opts.img_size), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])) train_loader = torch.utils.data.DataLoader(dataset, batch_size=opts.batch_size, shuffle=True, num_workers=opts.nThreads) # model print('\n--- load model ---') model = CDCGAN(opts) model.setgpu(opts.gpu) if opts.resume is None: model.initialize() ep0 = -1 total_it = 0 else: ep0, total_it = model.resume(opts.resume) ep0 += 1 print('start the training at epoch %d'%(ep0)) # saver for display and output saver = Saver(opts) # train print('\n--- train ---') max_it = 200000 for ep in range(ep0, opts.n_ep): for it, (images, label) in enumerate(train_loader): if images.size(0) != opts.batch_size: continue # input data images = images.cuda(opts.gpu).detach() # update model model.update_D(images, label) model.update_G() # save to display file if not opts.no_display_img: saver.write_display(total_it, model) print('total_it: %d (ep %d, it %d), lr %08f' % (total_it, ep, it, model.gen_opt.param_groups[0]['lr'])) total_it += 1 if total_it >= max_it: saver.write_img(-1, model) saver.write_model(-1, max_it, model) break # save result image saver.write_img(ep, model) # Save network weights saver.write_model(ep, total_it, model) return
def train_data(): opt = TrainOptions().parse() train_dataloader, test_dataloader = create_dataloader(opt) model = create_model(opt) optim = model.create_optimizer(opt) print('Start training.') width = len(str(opt.niter)) for it in range(1, opt.niter+1): it_str = '{:0{width}d}'.format(it, width=width) tqdm_desc = 'Epoch #{}/{:d}'.format(it_str, opt.niter) iter_start_time = time.time() train_pbar = tqdm(train_dataloader, desc=tqdm_desc) for data_i in train_pbar: metrics = train_one_step(model, optim, data_i) train_pbar.set_postfix(metrics) val_losses = [] val_accs = [] for data_i in test_dataloader: losses_sum, accs_sum = validate_one_step(model, data_i) val_losses.append(losses_sum) val_accs.append(accs_sum) val_size = len(test_dataloader.dataset) ave_val_loss = torch.stack(val_losses).sum().data.item() / val_size ave_val_acc = torch.stack(val_accs).sum().data.item() / val_size print('test loss: {:.5f}, test accuracy: {:.3%}, time elapsed {:.3f}s'.format( ave_val_loss, ave_val_acc, time.time() - iter_start_time )) if it % opt.save_epoch_freq == 0: save_network(model, it_str, opt)
def main(): parser = TrainOptions() args = parser.parse() args.mode = 'train' # Print CUDA version. print("Running code using CUDA {}".format(torch.version.cuda)) gpu_id = int(args.device[-1]) torch.cuda.set_device(gpu_id) print('Training on device cuda:{}'.format(gpu_id)) trainer = Trainer(args) if args.mode == 'train': trainer.train() elif args.mode == 'verify-data': trainer.verify_data()
def main(): from options import TrainOptions opt = TrainOptions().parse() dataloader = DataLoader(opt) dataloader_val = DataLoader_Val(opt) imgs_train, labels_train, infos_train = dataloader.read_inputs(opt.dataroot) imgs_val, labels_val, infos_val = dataloader_val.read_inputs(opt.dataroot_val) print(imgs_train, labels_train, infos_train,imgs_val,labels_val,infos_val)
def main(): log = logging.getLogger('pixsty') from options import TrainOptions parser = TrainOptions() parser.parser.add_argument('--subjects', type=str, nargs='+') args = parser.parse() log.info('Create dataset') train_loader = CustomDataLoader(args, phase='train') val_loader = CustomDataLoader(args, phase='val') print('training images = %d' % len(train_loader.dataset)) print('validation images = %d' % len(val_loader.dataset)) print('===> Build model') models = create_model(args) core_fn = { 'pix2pix': core.training_estimator, 'cyclegan': core.cyclegan_estimator, }[args.model] estimator_fn = core_fn(models, args) estimator_fn(train_loader, val_loader, epochs=args.niter)
def test(): from options import TrainOptions from nnf import NNFPredictor, Synthesiser3D opt = TrainOptions().parse() data_root = '/Users/tony/PycharmProjects/DynTexTrans/data/processed' dataset = DynTexFigureTrainDataset(data_root, 'flame') dataloader = DataLoader(dataset=dataset, batch_size=opt.batchsize, num_workers=opt.num_workers) syner = Synthesiser3D() nnfer = NNFPredictor() for i, (source_t, target_t, source_t1, target_t1) in enumerate(dataloader): print("-----------------") print(source_t.shape) print(target_t.shape) print(source_t1.shape) print(target_t1.shape) # source_t.requires_grad_() # target_t.requires_grad_() # source_t1.requires_grad_() # target_t1.requires_grad_() # # nnf = nnfer(source_t, target_t) name = os.path.join(data_root, '..', '{}.png'.format(str(i))) cv2.imwrite(name, (source_t.detach().numpy()[0].transpose(1, 2, 0) * 255).astype('int'))
def main(): opt = TrainOptions().parse() # writer writer = SummaryWriter(os.path.join(opt.log_dir, 'runs')) # dataset transform = transforms.Compose([ transforms.Resize((opt.input_size, opt.input_size)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ]) if opt.dataset == 'cifar10': dataset = torchvision.datasets.CIFAR10('data', train=True, download=True, transform=transform) elif opt.dataset == 'cifar100': dataset = torchvision.datasets.CIFAR100('data', train=True, download=True, transform=transform) elif opt.dataset == 'stl10': dataset = torchvision.datasets.STL10('data', split='train', download=True, transform=transform) loader = DataLoader(dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers) # model wgan_gp = WGAN_GP(opt) wgan_gp.train(loader, opt, writer)
def main(): # parse options parser = TrainOptions() opts = parser.parse() if opts.model_task != 'SYN' and opts.model_task != 'EDT': print('%s:unsupported task!'%opts.model_task) return SYN = opts.model_task == 'SYN' # create model print('--- create model ---') G_channels = 3 if SYN else 7 # SYN: (S), EDT:(S,I,M) D_channels = 6 if SYN else 7 # SYN: (S,I), EDT:(S,I,M) # (img_size, max_level) should be (256, 3), (128, 2) or (64, 1) netG = PSGAN(G_channels, opts.G_nlayers, opts.G_nf, D_channels, opts.D_nf, opts.D_nlayers, opts.max_dilate, opts.max_level, opts.img_size, opts.gpu!=0) if opts.gpu: netG.cuda() netG.init_networks(weights_init) netG.train() device = None if opts.gpu else torch.device('cpu') # to adapt G to the pretrained F if opts.use_F_level in [1,2,3]: # you could change the setting of netF based on your own pretraining setting netF_Norm = 'None' if SYN else 'BN' if opts.use_F_level == 1: netF = Pix2pix64(in_channels = G_channels, nef=64, useNorm=netF_Norm) elif opts.use_F_level == 2: netF = Pix2pix128(in_channels = G_channels, nef=64, useNorm=netF_Norm) else: netF = Pix2pix256(in_channels = G_channels, nef=64, useNorm=netF_Norm) netF.load_state_dict(torch.load(opts.load_F_name, map_location=device)) for param in netF.parameters(): param.requires_grad = False if opts.gpu: netF = netF.cuda() # for perceptual loss VGGNet = models.vgg19(pretrained=True).features VGGfeatures = VGGFeature(VGGNet, opts.gpu) for param in VGGfeatures.parameters(): param.requires_grad = False if opts.gpu: VGGfeatures.cuda() print('--- training ---') dataset = dset.ImageFolder(root=opts.train_path, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ])) # sampling refinement level l from l_sample_range # we found that r=0.5 should be trained to discriminate the results from those under r=0.0 l_sample_range = [1.*r/(opts.max_dilate-1) for r in range(opts.max_dilate)+[0.5]] # We train 64*64 image (max_level = 1) by directly training netG.G64 on 64*64. # We train 256*256 images (max_level = 3) by first training netG.G64 on 64*64, # then fixing netG.G64 and training netG.G128 on 128*128, # finally fixing netG.G128 and training netG.G256 on 256*256, like pix2pixHD # NOTE: using large batch size at level1,2 (small image resolution) could improve results # However, this could make CUDA out of memory when training goes to level3. # Because the memory used in level1,2 is not released. # Solution is to call train.py three times, # each time only training one level and saving the model parameters to load for the next level # for example, # 'for level in range(1,1+opts.max_level)' should be changed to 'for level in [cur_level]' # and when cur_level=2, adding 'netG.G64.load_state_dict(torch.load(saved model at level 1))' # when cur_level=3, adding 'netG.G128.load_state_dict(torch.load(saved model at level 2))' batchsize_level = [opts.batchsize_level1, opts.batchsize_level2, opts.batchsize_level3] # progressively training from level1 to max_level for level in range(1,1+opts.max_level): print('--- Training at level %d. Image resolution: %d ---' % (level, 2**(5+level))) # fix the model parameters if level in [2,3]: for param in netG.G64.parameters(): param.requires_grad = False if level in [3]: for param in netG.G128.parameters(): param.requires_grad = False dataloader = DataLoader(dataset, batch_size=batchsize_level[level-1], shuffle=True, num_workers=4, drop_last=True) print_step = int(len(dataloader) / 10.0) if not SYN: # generate random 2**(13-level) masks for sampling masks_num = 2**(13-level) all_masks = get_mask(masks_num, opts.img_size) all_masks = to_var(all_masks) if opts.gpu else all_masks # main iteration for epoch in range(opts.epoch_pre+opts.epoch): for i, data in enumerate(dataloader): # during pretraining epoches, we only train on the max refinement level l = 1.0 # then we will train on random l in [0,1] l = 1.0 if epoch < opts.epoch_pre else random.choice(l_sample_range) data = to_var(data[0]) if opts.gpu else data[0] # if input image is arranged as (S,I) use AtoB = True # if input image is arranged as (I,S) use AtoB = False if opts.AtoB: S = data[:,:,:,0:opts.img_size] I = data[:,:,:,opts.img_size:opts.img_size*2] else: S = data[:,:,:,opts.img_size:opts.img_size*2] I = data[:,:,:,0:opts.img_size] # apply netF loss, this will drastically increase the CUDA memory usuage netF_level = netF if level == opts.use_F_level else None if SYN: losses = netG.synthesis_one_pass(S, I, l, level, VGGfeatures, netF=netF_level) else: M = all_masks[torch.randint(masks_num, (batchsize_level[level-1],))] losses = netG.editing_one_pass(S, I, M, l, level, VGGfeatures, netF=netF_level) if i % print_step == 0: print('Epoch [%03d/%03d][%04d/%04d]' %(epoch+1, opts.epoch_pre+opts.epoch, i+1, len(dataloader)), end=': ') print('l: %+.3f, LD: %+.3f, LGadv: %+.3f, Lperc: %+.3f, Lrec: %+.3f'% (l, losses[0], losses[1], losses[2], losses[3])) print('--- Saving model at level %d ---' % level) netG.save_model(opts.save_model_path, opts.save_model_name, level=[level])
if __name__ == '__main__': import torch import torch.nn as nn import torchvision from torch.autograd import Variable import numpy as np from torch.nn import init import torch.optim as optim import math from math import log10 import time import datetime import cv2 from pytorch_msssim import ssim, ms_ssim from helper import * from model_full_no_stage2_4RIM import Deraining from data_with_grid import outdoor_rain_train, outdoor_rain_test from options import TrainOptions opt = TrainOptions().parse() torch.manual_seed(opt.seed) model = Deraining(opt) train_data = get_dataset(opt) test_data = get_testdataset(opt) train(opt, train_data, test_data, model)
def main(): # parse options parser = TrainOptions() opts = parser.parse() # daita loader print('\n--- load dataset ---') dataset = dataset_unpair(opts) train_loader = torch.utils.data.DataLoader(dataset, batch_size=opts.batch_size, shuffle=True, num_workers=opts.nThreads) # model print('\n--- load model ---') model = UID(opts) model.setgpu(opts.gpu) if opts.resume is None: model.initialize() ep0 = -1 total_it = 0 else: ep0, total_it = model.resume(opts.resume) model.set_scheduler(opts, last_ep=ep0) ep0 += 1 print('start the training at epoch %d' % (ep0)) # saver for display and output saver = Saver(opts) # train print('\n--- train ---') max_it = 500000 for ep in range(ep0, opts.n_ep): for it, (images_a, images_b) in enumerate(train_loader): if images_a.size(0) != opts.batch_size or images_b.size( 0) != opts.batch_size: continue images_a = images_a.cuda(opts.gpu).detach() images_b = images_b.cuda(opts.gpu).detach() # update model model.update_D(images_a, images_b) if (it + 1) % 2 != 0 and it != len(train_loader) - 1: continue model.update_EG() # save to display file if (it + 1) % 48 == 0: print('total_it: %d (ep %d, it %d), lr %08f' % (total_it + 1, ep, it + 1, model.gen_opt.param_groups[0]['lr'])) print( 'Dis_I_loss: %04f, Dis_B_loss %04f, GAN_loss_I %04f, GAN_loss_B %04f' % (model.disA_loss, model.disB_loss, model.gan_loss_i, model.gan_loss_b)) print('B_percp_loss %04f, Recon_II_loss %04f' % (model.B_percp_loss, model.l1_recon_II_loss)) if (it + 1) % 200 == 0: saver.write_img(ep * len(train_loader) + (it + 1), model) total_it += 1 if total_it >= max_it: saver.write_img(-1, model) saver.write_model(-1, model) break # decay learning rate if opts.n_ep_decay > -1: model.update_lr() # Save network weights saver.write_model(ep, total_it + 1, model) return
# PyTorch includes import torch from torch.autograd import Variable # Custom includes from options import TrainOptions from dataset import createDataset from models import createModel from visualizer import ProgressVisualizer assert Variable parser = TrainOptions() opt = parser.parse() # set dataloader trainDataset = createDataset(opt, split='train', nInput=opt.nInput) valDataset = createDataset(opt, split='val', nInput=opt.nInput) trainDataLoader = torch.utils.data.DataLoader(trainDataset, batch_size=opt.batchSize, shuffle=True, num_workers=opt.nThreads) valDataLoader = torch.utils.data.DataLoader(valDataset, batch_size=opt.batchSize, shuffle=True, num_workers=opt.nThreads) # set model model = createModel(opt) model.setup(opt)
def main(): opt = TrainOptions().parse() train(opt)
pbar = tqdm.tqdm(train_loader, desc=f'Epoch#{epoch}') for i, data in enumerate(pbar, 1): loss_terms, images = trainer.optimize_parameters( prepare_input(data), update_g=i % 5 == 0, update_d=True) pbar.set_postfix(history.add(loss_terms)) logger.image(images, epoch=epoch, prefix='train_') ''' validate ''' trainer.netG.eval() for data in val_loader: loss_terms, images = trainer.optimize_parameters( prepare_input(data), update_g=False, update_d=False) history.add(loss_terms, log_suffix='_val') logger.image(images, epoch=epoch, prefix='val_') logger.scalar(history.metric(), epoch) if epoch % opt.save_epoch_freq == 0: print(f'saving the model at the end of epoch {epoch}') trainer.save('latest') trainer.save(epoch) trainer.update_learning_rate() # clean the state of extensions history.clear() logger.clear() if __name__ == '__main__': parser = TrainOptions() parser.parser.add_argument('--subjects', type=str, nargs='+') main(parser.parse())
import time import sys import torch from torch.autograd import Variable from torchvision.utils import save_image from options import TrainOptions from models import Create_nets from datasets import Get_dataloader from optimizer import Get_loss_func, Get_optimizers from utils import ReplayBuffer, LambdaLR, sample_images #load the args args = TrainOptions().parse() # Calculate output of size discriminator (PatchGAN) patch = (1, args.img_height//(2**args.n_D_layers) - 2 , args.img_width//(2**args.n_D_layers) - 2) # Initialize generator and discriminator G__AB, D__B, G__BA, D__A = Create_nets(args) # Loss functions criterion_GAN, criterion_cycle, criterion_identity = Get_loss_func(args) # Optimizers optimizer_G, optimizer_D_B, optimizer_D_A = Get_optimizers(args, G__AB, G__BA, D__B, D__A ) # Learning rate update schedulers lr_scheduler_G = torch.optim.lr_scheduler.LambdaLR(optimizer_G, lr_lambda=LambdaLR(args.epoch_num, args.epoch_start, args.decay_epoch).step) lr_scheduler_D_B = torch.optim.lr_scheduler.LambdaLR(optimizer_D_B, lr_lambda=LambdaLR(args.epoch_num, args.epoch_start, args.decay_epoch).step) lr_scheduler_D_A = torch.optim.lr_scheduler.LambdaLR(optimizer_D_A, lr_lambda=LambdaLR(args.epoch_num, args.epoch_start, args.decay_epoch).step)
from torchvision import transforms import torchvision from options import TrainOptions import torch from torch.utils.data import DataLoader import os import torch.nn as nn import numpy as np from utils import * from LeNet_5 import LeNet_5 from time import localtime,strftime os.environ["CUDA_VISIBLE_DEVICES"] = '0' device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") opt, _ = TrainOptions() # prepare transform = transforms.Compose([ transforms.Grayscale(num_output_channels=1), transforms.Resize([32,32]), transforms.ToTensor() ]) train_data = ImageFolder(opt.train_data_path,transform = transform) train_loader = DataLoader(dataset=train_data, batch_size=opt.batch_size, shuffle=True) MyLeNet = LeNet_5() MyLeNet = MyLeNet.to(device) criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.Adam(MyLeNet.parameters(), lr=0.001)
print(':-o Goodbye') return False, None, None, None encoder, _ = factory.model_lc(args.encoder, args, device=trace.device) decoder, _ = factory.model_lc(args.decoder, args, device=trace.device) encoder.eval(), decoder.eval() choices = {0: to_exit, 1: sample, 2: interpolate, 3: hgmms, 4: last_try} menu = ' | '.join( sorted( list(f"{key}: {str(item).split()[1].split('.')[-1]}" for key, item in choices.items()))) eval_choice = 1 allow_saving = saving_handler(args, trace) while eval_choice: eval_choice = get_integer((0, len(choices)), menu + '\n') with torch.no_grad(): if eval_choice != len(choices) - 1: trace.memory = None check, trace.memory, image, points_group = choices[eval_choice]( encoder, decoder, args, trace) if check: function_name = str(trace.memory[0]).split()[1].split('.')[-1] allow_saving(function_name, image, points_group) print(f"{function_name} done") if __name__ == '__main__': cls = 'chair' evaluate(TrainOptions(tag=cls).load(), ViewMem())
def main(): # parse options parser = TrainOptions() opts = parser.parse() # daita loader print('\n--- load dataset ---') dataset = dataset_multi(opts) train_loader = torch.utils.data.DataLoader(dataset, batch_size=opts.batch_size, shuffle=True, num_workers=opts.nThreads) # model print('\n--- load model ---') model = SAVI2I(opts) model.setgpu(opts.gpu) if opts.resume is None: model.initialize() ep0 = -1 total_it = 0 else: ep0, total_it = model.resume(opts.resume) model.set_scheduler(opts, last_ep=ep0) ep0 += 1 print('start the training at epoch %d' % (ep0)) # saver for display and output saver = Saver(opts) # train print('\n--- train ---') max_it = 1000000 for ep in range(ep0, opts.n_ep): for it, (images, c_org, c_org_mask, c_org_id) in enumerate(train_loader): # input data images = torch.cat(images, dim=0) images = images.cuda(opts.gpu).detach() c_org = torch.cat(c_org, dim=0) c_org = c_org.cuda(opts.gpu).detach() c_org_mask = torch.cat(c_org_mask, dim=0) c_org_mask = c_org_mask.cuda(opts.gpu).detach() c_org_id = torch.cat(c_org_id, dim=0) c_org_id = c_org_id.cuda(opts.gpu).detach() # update model if (it + 1) % opts.d_iter != 0 and it < len(train_loader) - 2: model.update_D_content(images, c_org) continue else: model.update_D(images, c_org, c_org_mask, c_org_id) model.update_EFG() print('total_it: %d (ep %d, it %d), lr %08f' % (total_it, ep, it, model.gen_opt.param_groups[0]['lr'])) total_it += 1 if total_it >= max_it: saver.write_img(-1, model) saver.write_model(-1, max_it, model) break # decay learning rate if opts.n_ep_decay > -1: model.update_lr() # save result image saver.write_img(ep, model) # Save network weights saver.write_model(ep, total_it, model) return
def train_simple_trans(): opt = TrainOptions().parse() data_root = 'data/processed' train_params = {'lr': 0.01, 'epoch_milestones': (100, 500)} # dataset = DynTexNNFTrainDataset(data_root, 'flame') dataset = DynTexFigureTrainDataset(data_root, 'flame') dataloader = DataLoader(dataset=dataset, batch_size=opt.batchsize, num_workers=opt.num_workers, shuffle=True) nnf_conf = 3 syner = Synthesiser() nnfer = NNFPredictor(out_channel=nnf_conf) if torch.cuda.is_available(): syner = syner.cuda() nnfer = nnfer.cuda() optimizer_nnfer = Adam(nnfer.parameters(), lr=train_params['lr']) table = Table() for epoch in range(opt.epoch): pbar = tqdm(total=len(dataloader), desc='epoch#{}'.format(epoch)) pbar.set_postfix({'loss': 'N/A'}) loss_tot = 0.0 gamma = epoch / opt.epoch for i, (source_t, target_t, source_t1, target_t1) in enumerate(dataloader): if torch.cuda.is_available(): source_t = Variable(source_t, requires_grad=True).cuda() target_t = Variable(target_t, requires_grad=True).cuda() source_t1 = Variable(source_t1, requires_grad=True).cuda() target_t1 = Variable(target_t1, requires_grad=True).cuda() nnf = nnfer(source_t, target_t) if nnf_conf == 3: nnf = nnf[:, : 2, :, :] * nnf[:, 2:, :, :] # mask via the confidence # --- synthesis --- target_predict = syner(source_t, nnf) target_t1_predict = syner(source_t1, nnf) loss_t = tnf.mse_loss(target_predict, target_t) loss_t1 = tnf.mse_loss(target_t1_predict, target_t1) loss = loss_t + loss_t1 optimizer_nnfer.zero_grad() loss.backward() optimizer_nnfer.step() loss_tot += float(loss_t1) # --- vis --- name = os.path.join(data_root, '../result/', str(epoch), '{}.png'.format(str(i))) index = str(epoch) + '({})'.format(i) if not os.path.exists('/'.join(name.split('/')[:-1])): os.makedirs('/'.join(name.split('/')[:-1])) cv2.imwrite( name.replace('.png', '_s.png'), (source_t.detach().cpu().numpy()[0].transpose(1, 2, 0) * 255).astype('int')) cv2.imwrite( name.replace('.png', '_s1.png'), (source_t1.detach().cpu().numpy()[0].transpose(1, 2, 0) * 255).astype('int')) cv2.imwrite( name.replace('.png', '_t.png'), (target_t.detach().cpu().numpy()[0].transpose(1, 2, 0) * 255).astype('int')) cv2.imwrite( name.replace('.png', '_p.png'), (target_predict.detach().cpu().numpy()[0].transpose(1, 2, 0) * 255).astype('int')) cv2.imwrite( name.replace('.png', '_t1.png'), (target_t1.detach().cpu().numpy()[0].transpose(1, 2, 0) * 255).astype('int')) cv2.imwrite(name.replace('.png', '_p1.png'), (target_t1_predict.detach().cpu().numpy()[0].transpose( 1, 2, 0) * 255).astype('int')) # vis in table table.add( index, os.path.abspath(name.replace('.png', '_s.png')).replace( '/mnt/cephfs_hl/lab_ad_idea/maoyiming', '')) table.add( index, os.path.abspath(name.replace('.png', '_s1.png')).replace( '/mnt/cephfs_hl/lab_ad_idea/maoyiming', '')) table.add( index, os.path.abspath(name.replace('.png', '_t.png')).replace( '/mnt/cephfs_hl/lab_ad_idea/maoyiming', '')) table.add( index, os.path.abspath(name.replace('.png', '_t1.png')).replace( '/mnt/cephfs_hl/lab_ad_idea/maoyiming', '')) table.add( index, os.path.abspath(name.replace('.png', '_p.png')).replace( '/mnt/cephfs_hl/lab_ad_idea/maoyiming', '')) table.add( index, os.path.abspath(name.replace('.png', '_p1.png')).replace( '/mnt/cephfs_hl/lab_ad_idea/maoyiming', '')) pbar.set_postfix({'loss': str(loss_tot / (i + 1))}) pbar.update(1) table.build_html('data/') pbar.close()
def train_complex_trans(): opt = TrainOptions().parse() data_root = 'data/processed' train_params = {'lr': 0.001, 'epoch_milestones': (100, 500)} dataset = DynTexFigureTransTrainDataset(data_root, 'flame') dataloader = DataLoader(dataset=dataset, batch_size=opt.batchsize, num_workers=opt.num_workers, shuffle=True) nnf_conf = 3 syner = Synthesiser() nnfer = NNFPredictor(out_channel=nnf_conf) flownet = NNFPredictor(out_channel=nnf_conf) if torch.cuda.is_available(): syner = syner.cuda() nnfer = nnfer.cuda() flownet = flownet.cuda() optimizer_nnfer = Adam(nnfer.parameters(), lr=train_params['lr']) optimizer_flow = Adam(flownet.parameters(), lr=train_params['lr'] * 0.1) scheduler_nnfer = lr_scheduler.MultiStepLR( optimizer_nnfer, gamma=0.1, last_epoch=-1, milestones=train_params['epoch_milestones']) scheduler_flow = lr_scheduler.MultiStepLR( optimizer_flow, gamma=0.1, last_epoch=-1, milestones=train_params['epoch_milestones']) table = Table() writer = SummaryWriter(log_dir=opt.log_dir) for epoch in range(opt.epoch): scheduler_flow.step() scheduler_nnfer.step() pbar = tqdm(total=len(dataloader), desc='epoch#{}'.format(epoch)) pbar.set_postfix({'loss': 'N/A'}) loss_tot = 0.0 for i, (source_t, target_t, source_t1, target_t1) in enumerate(dataloader): if torch.cuda.is_available(): source_t = Variable(source_t, requires_grad=True).cuda() target_t = Variable(target_t, requires_grad=True).cuda() source_t1 = Variable(source_t1, requires_grad=True).cuda() target_t1 = Variable(target_t1, requires_grad=True).cuda() nnf = nnfer(source_t, target_t) flow = flownet(source_t, source_t1) # mask... if nnf_conf == 3: nnf = nnf[:, : 2, :, :] * nnf[:, 2:, :, :] # mask via the confidence flow = flow[:, :2, :, :] * flow[:, 2:, :, :] # --- synthesis --- source_t1_predict = syner(source_t, flow) # flow penalty target_flow = syner(flow, nnf) # predict flow target_t1_predict = syner(target_t, target_flow) # target_t1_predict = syner(source_t1, nnf) loss_t1_f = tnf.mse_loss(source_t1, source_t1_predict) # flow penalty loss_t1 = tnf.mse_loss(target_t1_predict, target_t1) # total penalty loss = loss_t1_f + loss_t1 * 2 optimizer_flow.zero_grad() optimizer_nnfer.zero_grad() loss.backward() optimizer_nnfer.step() optimizer_flow.step() loss_tot += float(loss) # --- vis --- if epoch % 10 == 0 and i % 2 == 0: name = os.path.join(data_root, '../result/', str(epoch), '{}.png'.format(str(i))) index = str(epoch) + '({})'.format(i) if not os.path.exists('/'.join(name.split('/')[:-1])): os.makedirs('/'.join(name.split('/')[:-1])) cv2.imwrite( name.replace('.png', '_s.png'), (source_t.detach().cpu().numpy()[0].transpose(1, 2, 0) * 255).astype('int')) cv2.imwrite( name.replace('.png', '_s1.png'), (source_t1.detach().cpu().numpy()[0].transpose(1, 2, 0) * 255).astype('int')) cv2.imwrite( name.replace('.png', '_t.png'), (target_t.detach().cpu().numpy()[0].transpose(1, 2, 0) * 255).astype('int')) cv2.imwrite( name.replace('.png', '_p.png'), (source_t1_predict.detach().cpu().numpy()[0].transpose( 1, 2, 0) * 255).astype('int')) cv2.imwrite( name.replace('.png', '_t1.png'), (target_t1.detach().cpu().numpy()[0].transpose(1, 2, 0) * 255).astype('int')) cv2.imwrite( name.replace('.png', '_p1.png'), (target_t1_predict.detach().cpu().numpy()[0].transpose( 1, 2, 0) * 255).astype('int')) # vis in table table.add( index, os.path.abspath(name.replace('.png', '_s.png')).replace( '/mnt/cephfs_hl/lab_ad_idea/maoyiming', '')) table.add( index, os.path.abspath(name.replace('.png', '_s1.png')).replace( '/mnt/cephfs_hl/lab_ad_idea/maoyiming', '')) table.add( index, os.path.abspath(name.replace('.png', '_t.png')).replace( '/mnt/cephfs_hl/lab_ad_idea/maoyiming', '')) table.add( index, os.path.abspath(name.replace('.png', '_t1.png')).replace( '/mnt/cephfs_hl/lab_ad_idea/maoyiming', '')) table.add( index, os.path.abspath(name.replace('.png', '_p.png')).replace( '/mnt/cephfs_hl/lab_ad_idea/maoyiming', '')) table.add( index, os.path.abspath(name.replace('.png', '_p1.png')).replace( '/mnt/cephfs_hl/lab_ad_idea/maoyiming', '')) pbar.set_postfix({'loss': str(loss_tot / (i + 1))}) writer.add_scalar('scalars/{}/loss_train'.format(opt.time), float(loss), i + int(epoch * len(dataloader))) writer.add_scalar('scalars/{}/lr'.format(opt.time), float(scheduler_nnfer.get_lr()[0]), i + int(epoch * len(dataloader))) pbar.update(1) table.build_html('data/') pbar.close() writer.close()
import data import models import optimizers from options import TrainOptions from util import IterationCounter from util import Visualizer from util import MetricTracker from evaluation import GroupEvaluator opt = TrainOptions().parse() dataset = data.create_dataset(opt) opt.dataset = dataset iter_counter = IterationCounter(opt) visualizer = Visualizer(opt) metric_tracker = MetricTracker(opt) evaluators = GroupEvaluator(opt) model = models.create_model(opt) optimizer = optimizers.create_optimizer(opt, model) while not iter_counter.completed_training(): with iter_counter.time_measurement("data"): cur_data = next(dataset) with iter_counter.time_measurement("train"): losses = optimizer.train_one_step(cur_data, iter_counter.steps_so_far) metric_tracker.update_metrics(losses, smoothe=True) with iter_counter.time_measurement("maintenance"): if iter_counter.needs_printing(): visualizer.print_current_losses(iter_counter.steps_so_far,
def main(): # parse options parser = TrainOptions() opts = parser.parse() # daita loader print('\n--- load dataset ---') vocab = pickle.load( open(os.path.join(opts.vocab_path, '%s_vocab.pkl' % opts.data_name), 'rb')) vocab_size = len(vocab) opts.vocab_size = vocab_size torch.backends.cudnn.enabled = False # Load data loaders train_loader, val_loader = data.get_loaders(opts.data_name, vocab, opts.crop_size, opts.batch_size, opts.workers, opts) test_loader = data.get_test_loader('test', opts.data_name, vocab, opts.crop_size, opts.batch_size, opts.workers, opts) # model print('\n--- load subspace ---') subspace = model_2.VSE(opts) subspace.setgpu() print('\n--- load model ---') model = DRIT(opts) model.setgpu(opts.gpu) if opts.resume is None: #之前没有保存过模型 model.initialize() ep0 = -1 total_it = 0 else: ep0, total_it = model.resume(opts.resume) model.set_scheduler(opts, last_ep=ep0) ep0 += 1 print('start the training at epoch %d' % (ep0)) # saver for display and output saver = Saver(opts) # train print('\n--- train ---') max_it = 500000 score = 0.0 subspace.train_start() for ep in range(ep0, opts.pre_iter): print('-----ep:{} --------'.format(ep)) for it, (images, captions, lengths, ids) in enumerate(train_loader): if it >= opts.train_iter: break # input data images = images.cuda(opts.gpu).detach() captions = captions.cuda(opts.gpu).detach() img, cap = subspace.train_emb(images, captions, lengths, ids, pre=True) #[b,1024] subspace.pre_optimizer.zero_grad() img = img.view(images.size(0), -1, 32, 32) cap = cap.view(images.size(0), -1, 32, 32) model.pretrain_ae(img, cap) if opts.grad_clip > 0: clip_grad_norm(subspace.params, opts.grad_clip) subspace.pre_optimizer.step() for ep in range(ep0, opts.n_ep): subspace.train_start() adjust_learning_rate(opts, subspace.optimizer, ep) for it, (images, captions, lengths, ids) in enumerate(train_loader): if it >= opts.train_iter: break # input data images = images.cuda(opts.gpu).detach() captions = captions.cuda(opts.gpu).detach() img, cap = subspace.train_emb(images, captions, lengths, ids) #[b,1024] img = img.view(images.size(0), -1, 32, 32) cap = cap.view(images.size(0), -1, 32, 32) subspace.optimizer.zero_grad() for p in model.disA.parameters(): p.requires_grad = True for p in model.disB.parameters(): p.requires_grad = True for p in model.disA_attr.parameters(): p.requires_grad = True for p in model.disB_attr.parameters(): p.requires_grad = True for i in range(opts.niters_gan_d): #5 model.update_D(img, cap) for p in model.disA.parameters(): p.requires_grad = False for p in model.disB.parameters(): p.requires_grad = False for p in model.disA_attr.parameters(): p.requires_grad = False for p in model.disB_attr.parameters(): p.requires_grad = False for i in range(opts.niters_gan_enc): model.update_E(img, cap) #利用新的content损失函数 subspace.optimizer.step() print('total_it: %d (ep %d, it %d), lr %09f' % (total_it, ep, it, model.gen_opt.param_groups[0]['lr'])) total_it += 1 # decay learning rate if opts.n_ep_decay > -1: model.update_lr() # save result image #saver.write_img(ep, model) if (ep + 1) % opts.n_ep == 0: print('save model') filename = os.path.join(opts.result_dir, opts.name) model.save('%s/final_model.pth' % (filename), ep, total_it) torch.save(subspace.state_dict(), '%s/final_subspace.pth' % (filename)) elif (ep + 1) % 10 == 0: print('save model') filename = os.path.join(opts.result_dir, opts.name) model.save('%s/%s_model.pth' % (filename, str(ep + 1)), ep, total_it) torch.save(subspace.state_dict(), '%s/%s_subspace.pth' % (filename, str(ep + 1))) if (ep + 1) % opts.model_save_freq == 0: a = None b = None c = None d = None subspace.val_start() for it, (images, captions, lengths, ids) in enumerate(test_loader): if it >= opts.val_iter: break images = images.cuda(opts.gpu).detach() captions = captions.cuda(opts.gpu).detach() img_emb, cap_emb = subspace.forward_emb(images, captions, lengths, volatile=True) img = img_emb.view(images.size(0), -1, 32, 32) cap = cap_emb.view(images.size(0), -1, 32, 32) image1, text1 = model.test_model2(img, cap) img2 = image1.view(images.size(0), -1) cap2 = text1.view(images.size(0), -1) if a is None: a = np.zeros( (opts.val_iter * opts.batch_size, img_emb.size(1))) b = np.zeros( (opts.val_iter * opts.batch_size, cap_emb.size(1))) c = np.zeros( (opts.val_iter * opts.batch_size, img2.size(1))) d = np.zeros( (opts.val_iter * opts.batch_size, cap2.size(1))) a[ids] = img_emb.data.cpu().numpy().copy() b[ids] = cap_emb.data.cpu().numpy().copy() c[ids] = img2.data.cpu().numpy().copy() d[ids] = cap2.data.cpu().numpy().copy() aa = torch.from_numpy(a) bb = torch.from_numpy(b) cc = torch.from_numpy(c) dd = torch.from_numpy(d) (r1, r5, r10, medr, meanr) = i2t(aa, bb, measure=opts.measure) print('test640: subspace: Med:{}, r1:{}, r5:{}, r10:{}'.format( medr, r1, r5, r10)) (r1i, r5i, r10i, medri, meanr) = t2i(aa, bb, measure=opts.measure) print('test640: subspace: Med:{}, r1:{}, r5:{}, r10:{}'.format( medri, r1i, r5i, r10i)) (r2, r3, r4, m1, m2) = i2t(cc, dd, measure=opts.measure) print('test640: encoder: Med:{}, r1:{}, r5:{}, r10:{}'.format( m1, r2, r3, r4)) (r2i, r3i, r4i, m1i, m2i) = t2i(cc, dd, measure=opts.measure) print('test640: encoder: Med:{}, r1:{}, r5:{}, r10:{}'.format( m1i, r2i, r3i, r4i)) curr = r2 + r3 + r4 + r2i + r3i + r4i if curr > score: score = curr print('save model') filename = os.path.join(opts.result_dir, opts.name) model.save('%s/best_model.pth' % (filename), ep, total_it) torch.save(subspace.state_dict(), '%s/subspace.pth' % (filename)) a = None b = None c = None d = None for it, (images, captions, lengths, ids) in enumerate(test_loader): images = images.cuda(opts.gpu).detach() captions = captions.cuda(opts.gpu).detach() img_emb, cap_emb = subspace.forward_emb(images, captions, lengths, volatile=True) img = img_emb.view(images.size(0), -1, 32, 32) cap = cap_emb.view(images.size(0), -1, 32, 32) image1, text1 = model.test_model2(img, cap) img2 = image1.view(images.size(0), -1) cap2 = text1.view(images.size(0), -1) if a is None: a = np.zeros((len(test_loader.dataset), img_emb.size(1))) b = np.zeros((len(test_loader.dataset), cap_emb.size(1))) c = np.zeros((len(test_loader.dataset), img2.size(1))) d = np.zeros((len(test_loader.dataset), cap2.size(1))) a[ids] = img_emb.data.cpu().numpy().copy() b[ids] = cap_emb.data.cpu().numpy().copy() c[ids] = img2.data.cpu().numpy().copy() d[ids] = cap2.data.cpu().numpy().copy() aa = torch.from_numpy(a) bb = torch.from_numpy(b) cc = torch.from_numpy(c) dd = torch.from_numpy(d) (r1, r5, r10, medr, meanr) = i2t(aa, bb, measure=opts.measure) print('test5000: subspace: Med:{}, r1:{}, r5:{}, r10:{}'.format( medr, r1, r5, r10)) (r1i, r5i, r10i, medri, meanr) = t2i(aa, bb, measure=opts.measure) print('test5000: subspace: Med:{}, r1:{}, r5:{}, r10:{}'.format( medri, r1i, r5i, r10i)) (r2, r3, r4, m1, m2) = i2t(cc, dd, measure=opts.measure) print('test5000: encoder: Med:{}, r1:{}, r5:{}, r10:{}'.format( m1, r2, r3, r4)) (r2i, r3i, r4i, m1i, m2i) = t2i(cc, dd, measure=opts.measure) print('test5000: encoder: Med:{}, r1:{}, r5:{}, r10:{}'.format( m1i, r2i, r3i, r4i)) return
def done(): from deprecated.kpdetector import KeyPointDetector, StackedHourglass from options import TrainOptions from dataloader import DynTexTrainDataset from torch.utils.data import DataLoader from torch.optim.lr_scheduler import MultiStepLR from tqdm import tqdm from pprint import pprint import cv2 opt = TrainOptions().parse() print('---- training params. ----') train_params = {'lr': 0.001, 'epoch_milestones': (100, 500)} loss_weight = { 'reconstruction_def': 1.0, 'reconstruction': 10.0, 'generator_gan': 1.0, 'discriminator_gan': 1.0 } pprint(train_params) pprint(loss_weight) kpd = KeyPointDetector(opt).train() gen = Generator(opt).train() dis_ = Discriminator(opt).train() gen_vanila = StackedHourglass(in_channels=6, out_channels=3, inter_channels=128, stacked_num=2, dropout_rate=0.1, refine=True) start_epoch = 0 optimizer_gv = torch.optim.Adam(gen_vanila.parameters(), lr=train_params['lr'], betas=(0.5, 0.999)) scheduler_gv = MultiStepLR(optimizer_gv, train_params['epoch_milestones'], gamma=0.1, last_epoch=start_epoch - 1) optimizer_generator = torch.optim.Adam(gen.parameters(), lr=train_params['lr'], betas=(0.5, 0.999)) optimizer_discriminator = torch.optim.Adam(dis_.parameters(), lr=train_params['lr'], betas=(0.5, 0.999)) optimizer_kp_detector = torch.optim.Adam(kpd.parameters(), lr=train_params['lr'], betas=(0.5, 0.999)) scheduler_generator = MultiStepLR(optimizer_generator, train_params['epoch_milestones'], gamma=0.1, last_epoch=start_epoch - 1) scheduler_discriminator = MultiStepLR(optimizer_discriminator, train_params['epoch_milestones'], gamma=0.1, last_epoch=start_epoch - 1) scheduler_kp_detector = MultiStepLR(optimizer_kp_detector, train_params['epoch_milestones'], gamma=0.1, last_epoch=start_epoch - 1) data_root = '/Users/tony/PycharmProjects/DynTexTrans/data/processed' dataset = DynTexTrainDataset(data_root, 'flame') dataloader = DataLoader(dataset=dataset, batch_size=opt.batchsize, num_workers=opt.num_workers, shuffle=True) for epoch in range(start_epoch, start_epoch + 700): pbar = tqdm(total=len(dataloader), desc="=> training epoch # {}".format(epoch), ascii=True) pbar.set_postfix({'L_G': 'N/A', 'L_D': 'N/A'}) for batch, (source_tensor, drivin_tensor) in enumerate(dataloader): source_tensor = torch.Tensor.float(source_tensor).requires_grad_() drivin_tensor = torch.Tensor.float(drivin_tensor).requires_grad_() inputs = torch.cat([source_tensor, drivin_tensor], 1) outs = gen_vanila(inputs) print(outs[-1].shape, drivin_tensor.shape) gened = outs[-1] # source_kp, drivin_kp = kpd(torch.cat([source_tensor, drivin_tensor], dim=1)) # gened, inter = gen(source_tensor, drivin_kp, source_kp) # dis_feat = dis_(gened) # dis_real = dis_(drivin_tensor) # loss_g = generator_loss(dis_feat, dis_real, loss_weights=loss_weight) # loss_g = torch.mean(torch.cat(list(loss_g.values())), dim=0) # loss_g.backward(retain_graph=True) mse = torch.mean((gened - drivin_tensor) * (gened - drivin_tensor)) mse.backward(retain_graph=True) optimizer_gv.step() optimizer_gv.zero_grad() # optimizer_generator.step() # optimizer_generator.zero_grad() # optimizer_discriminator.zero_grad() # loss_d = discriminator_loss(dis_feat, dis_real, weight=loss_weight['discriminator_gan']) # loss_d = torch.mean(torch.cat(list(loss_d.values())), dim=0) # loss_d.backward() # optimizer_discriminator.step() # optimizer_discriminator.zero_grad() # # optimizer_kp_detector.step() # optimizer_kp_detector.zero_grad() # vis ge = normalization(gened.detach().numpy()[0].transpose( [1, 2, 0])).astype('uint8') # ge = (inter.detach().numpy()[0]) # print(len(np.unique(ge)), np.max(np.unique(ge)), np.min(np.unique(ge)), np.mean(ge)) gt = (drivin_tensor.detach().numpy()[0].transpose([1, 2, 0]) * 255).astype('uint8') if batch == 1: root = '/Users/tony/PycharmProjects/DynTexTrans/data/result/' vis = np.zeros(shape=(*ge.shape[:2], 3)) vis[:, :, :ge.shape[2]] = ge cv2.imwrite(root + '{}_{}_ge.png'.format(epoch, batch), vis) cv2.imwrite(root + 'gt.png', gt) # vis # loss_dict = {'L_G': '{:.5f}'.format(float(loss_g)), 'L_D': '{:.5f}'.format(float(loss_d))} # pbar.set_postfix(loss_dict) loss = {'loss': '{:.5f}'.format(float(mse))} pbar.set_postfix(loss) pbar.update(1) scheduler_discriminator.step(epoch) scheduler_generator.step(epoch) scheduler_kp_detector.step(epoch) pbar.close()
def main(): global val_accuracy, val_f1, val_precision, val_recall, val_cross_corr_a, val_cross_corr_b, val_mse, val_mae, \ saver, ep, save_opts, total_it, iter_counter, t0 # initialise params parser = TrainOptions() opts = parser.parse() opts.random_seed = RANDOM_SEED opts.device = opts.device if torch.cuda.is_available( ) and opts.gpu else 'cpu' opts.name = opts.data_type + '_' + time.strftime("%d%m%Y-%H%M") opts.results_path = os.path.join(opts.result_dir, opts.name) opts.image_size = IMAGE_SIZE opts.age_range_0 = AGE_RANGE_0 opts.age_range_1 = AGE_RANGE_1 opts.resize_image = RESIZE_IMAGE opts.resize_size = RESIZE_SIZE ep0 = 0 total_it = 0 val_accuracy = np.zeros(opts.n_ep) val_f1 = np.zeros(opts.n_ep) val_precision = np.zeros(opts.n_ep) val_recall = np.zeros(opts.n_ep) val_cross_corr_a = np.zeros(opts.n_ep) val_cross_corr_b = np.zeros(opts.n_ep) val_mse = np.zeros(opts.n_ep) val_mae = np.zeros(opts.n_ep) t0 = time.time() # saver for display and output if opts.data_dim == '3d': from saver_3d import Saver opts.nz = LATENT_3D else: from saver import Saver opts.nz = LATENT_2D print('\n--- load dataset ---') # add new dataloader in _load_dataloader(), and in dataloader_utils.py healthy_dataloader, healthy_val_dataloader, healthy_test_dataloader, \ anomaly_dataloader, anomaly_val_dataloader, anomaly_test_dataloader = _load_dataloader(opts) print('\n--- load model ---') model = ICAM(opts) model.setgpu(opts.device) model.initialize() model.set_scheduler(opts, last_ep=ep0) save_opts = vars(opts) saver = Saver(opts) if not os.path.exists(opts.results_path): os.makedirs(opts.results_path) with open(opts.results_path + '/parameters.json', 'w') as file: json.dump(save_opts, file, indent=4, sort_keys=True) print('\n--- train ---') for ep in range(ep0, opts.n_ep): healthy_data_iter = iter(healthy_dataloader) anomaly_data_iter = iter(anomaly_dataloader) iter_counter = 0 while iter_counter < len(anomaly_dataloader) and iter_counter < len( healthy_dataloader): # output of iter dataloader: [tensor image, tensor label (regression), tensor mask] healthy_images, healthy_label_reg, healthy_mask = healthy_data_iter.next( ) anomaly_images, anomaly_label_reg, anomaly_mask = anomaly_data_iter.next( ) healthy_c_org = torch.zeros( (healthy_images.size(0), opts.num_domains)).to(opts.device) healthy_c_org[:, 0] = 1 anomaly_c_org = torch.zeros( (healthy_images.size(0), opts.num_domains)).to(opts.device) anomaly_c_org[:, 1] = 1 images = torch.cat((healthy_images, anomaly_images), dim=0).type(torch.FloatTensor) c_org = torch.cat((healthy_c_org, anomaly_c_org), dim=0).type(torch.FloatTensor) label_reg = torch.cat((healthy_label_reg, anomaly_label_reg), dim=0).type(torch.FloatTensor) if len(healthy_mask.size()) > 2: mask = torch.cat((healthy_mask, anomaly_mask), dim=0).type(torch.FloatTensor) mask = mask.to(opts.device).detach() else: mask = None iter_counter += 1 if images.size(0) != opts.batch_size: continue # input data images = images.to(opts.device).detach() c_org = c_org.to(opts.device).detach() label_reg = label_reg.to(opts.device).detach() # update model if (iter_counter % opts.d_iter) != 0 and iter_counter < len( anomaly_dataloader) - opts.d_iter: model.update_D_content(opts, images, c_org) continue model.update_D(opts, images, c_org, label_reg, mask=mask) model.update_EG(opts) if ((total_it + 1) % opts.train_print_it) == 0: train_accuracy, train_f1, _, _ = model.classification_scores( images, c_org) if opts.regression: train_mse, train_mae, _ = model.regression( images, label_reg) if total_it == 0: saver.write_img(ep, total_it, model) elif total_it % opts.display_freq == 0: saver.write_img(ep, total_it, model) total_it += 1 # save to tensorboard saver.write_display(total_it, model) time_elapsed = time.time() - t0 hours, rem = divmod(time_elapsed, 3600) minutes, seconds = divmod(rem, 60) if (total_it % opts.train_print_it) == 0: print( 'Total it: {:d} (ep {:d}, it {:d}), Accuracy: {:.2f}, F1 score: {:.2f}, ' 'Elapsed time: {:0>2}:{:0>2}:{:05.2f}'.format( total_it, ep, iter_counter, train_accuracy, train_f1, int(hours), int(minutes), seconds)) # save model if ep % opts.model_save_freq == 0: saver.write_model(ep, total_it, 0, model, epoch=True) saver.write_img(ep, total_it, model) # example validation try: _validation(opts, model, healthy_val_dataloader, anomaly_val_dataloader) except Exception as e: print(f'Encountered error during validation - {e}') raise e # example test try: _test(opts, model, healthy_test_dataloader, anomaly_test_dataloader) except Exception as e: print(f'Encountered error during validation - {e}') raise e # save last model saver.write_model(ep, total_it, iter_counter, model, model_name='model_last') saver.write_img(ep, total_it, model) return
def main(): # for multi processing if torch.cuda.is_available(): multiprocessing.set_start_method('spawn') # parse options parser = TrainOptions() opts = parser.parse() if (z_delete_mode): z_zero_train = torch.zeros(opts.batch_size, 32).cuda(opts.gpu) z_zero_test = torch.zeros(opts.batch_size_test, 32).cuda(opts.gpu) # daita loader print('\n--- load dataset ---') sharp_whole_root = None if (sharp_whole_mode): sharp_whole_root = "../datasets_npz/datasets_sharp_whole_" + data_version dataset = dataset_pair_group_simple( opts, "../datasets_npz/datasets_blur_" + data_version, "../datasets_npz/datasets_sharp_" + data_version, "../datasets_npz/datasets_sharp_start_" + data_version, "../datasets_npz/datasets_gt_vel_" + data_version, "../datasets_npz/datasets_gt_pos_" + data_version, "../datasets_npz/datasets_shutter_speed_" + data_version, sharp_whole_root=sharp_whole_root) train_loader = torch.utils.data.DataLoader(dataset, batch_size=opts.batch_size, shuffle=True, num_workers=opts.nThreads) dataset_test = dataset_pair_group_simple( opts, "../datasets_npz/datasets_blur_" + data_version, "../datasets_npz/datasets_sharp_" + data_version, "../datasets_npz/datasets_sharp_start_" + data_version, "../datasets_npz/datasets_gt_vel_" + data_version, "../datasets_npz/datasets_gt_pos_" + data_version, "../datasets_npz/datasets_shutter_speed_" + data_version, test_mode=True, sharp_whole_root=sharp_whole_root) test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=opts.batch_size_test, shuffle=False, num_workers=opts.nThreads) test_iter = iter(test_loader) # model print('\n--- load model ---') model = UID(opts) model.setgpu(opts.gpu) if opts.resume is None: model.initialize() ep0 = -1 total_it = 0 else: ep0, total_it = model.resume(opts.resume) model.set_scheduler(opts, last_ep=ep0) ep0 += 1 print('start the training at epoch %d' % (ep0)) # saver for display and output saver = Saver(opts) # train print('\n--- train ---') max_it = 500000 for ep in range(ep0, opts.n_ep): for it, (images_b_s, gt_s, gt_pos_s, shutter_speed_s, images_s_whole_s, images_a_s) in enumerate(train_loader): if images_b_s.size(0) != opts.batch_size or images_b_s.size( 0) != opts.batch_size: continue images_b_s = images_b_s.cuda(opts.gpu).detach() for o in range(n_obj): gt_s[o] = gt_s[o].cuda(opts.gpu).detach() gt_pos_s[o] = gt_pos_s[o].cuda(opts.gpu).detach() shutter_speed_s = shutter_speed_s.cuda(opts.gpu).detach() if (sharp_whole_mode): images_s_whole_s = images_s_whole_s.cuda(opts.gpu).detach() else: images_a_s = images_a_s.cuda(opts.gpu).detach() # update model s = 0 z = None #last_gt = gt_s[:,opts.sequence_num - 1,:] last_gt = None loss_pred_last_sum = 0 while (s < (opts.sequence_num - 2)): if (z_delete_mode): z = z_zero_train # now: s+1, max: sequence_num-1 left_steps = opts.sequence_num - s - 2 # 현재(s+1)스텝으로부터 몇번까지 더 gt 가 있는지 images_prev = images_b_s[:, s, :, :, :] images_b = images_b_s[:, s + 1, :, :, :] images_post = images_b_s[:, s + 2, :, :, :] gt = [] gt_pos = [] for o in range(n_obj): gt.append(gt_s[o][:, s + 1, :]) gt_pos.append(gt_pos_s[o][:, s + 1, :]) next_gt = [] next_gt_pos = [] for o in range(n_obj): next_gt.append(gt_s[o][:, s + 2, :]) next_gt_pos.append(gt_pos_s[o][:, s + 2, :]) shutter_speed = shutter_speed_s if (sharp_whole_mode): images_s_whole = images_s_whole_s[:, (s + 1) * 16:(s + 2) * 16, :, :, :] images_a = None else: images_s_whole = None images_a = images_a_s[:, s + 1, :, :, :] if throwing_mode and s != 0: given_vel = model.next_vel_pred given_pos = model.next_pos_pred else: given_vel = None given_pos = None model.update(images_prev, images_post, images_b, gt, gt_pos, next_gt=next_gt, next_gt_pos=next_gt_pos, z=z, last_gt=last_gt, left_steps=left_steps, gt_pos_set=gt_pos_s, shutter_speed=shutter_speed, images_s_whole=images_s_whole, images_a=images_a, given_vel=given_vel, given_pos=given_pos) s += 1 if (model.pred_mode): z = model.z_next loss_pred_last_sum += model.loss_pred_last_vel else: z = None loss_pred_last_sum += -1 # save to display file if (it + 1) % 1 == 0: print('total_it: %d (ep %d, it %d), lr %08f' % (total_it + 1, ep, it + 1, model.enc_c_opt.param_groups[0]['lr'])) if pred_mode: loss_pred_vel = model.loss_pred_vel[0] else: loss_pred_vel = -1 print( 'gen_loss: %04f, vel_loss %04f, vel_recons_loss %04f, gen_loss_gt_vel %04f, inverse_loss %04f, vel_dir_loss %04f, loss_content %04f, vel_pred_loss %04f, pos_loss %04f, last_vel_pred_loss %04f, whole_pos_pred %04f, loss_sharp %04f' % (model.loss_gen, model.loss_vel[0], model.loss_vel_recons, model.loss_gen_gt_vel, model.loss_inverse, model.loss_vel_dir[0], model.loss_content, loss_pred_vel, model.loss_pos[0], loss_pred_last_sum / s, model.loss_sharp_whole, model.loss_sharp)) #print(model.next_gt) #print(model.next_gt_pred) if (it + 1) % test_ratio == 0: n = 0 total_dir_loss = 0 total_speed_loss = 0 total_vel_loss = 0 total_pos_loss = 0 total_dir_entropy = 0 for rep in range(1): try: images_b_s, gt_s, gt_pos_s, shutter_speed_s, images_s_whole_s, images_a_s = next( test_iter) except StopIteration: test_iter = iter(test_loader) images_b_s, gt_s, gt_pos_s, shutter_speed_s, _, _ = next( test_iter) while (images_b_s.size(0) != opts.batch_size_test): try: images_b_s, gt_s, gt_pos_s, shutter_speed_s, _, _ = next( test_iter) except StopIteration: test_iter = iter(test_loader) images_b_s, gt_s, gt_pos_s, shutter_speed_s, _, _ = next( test_iter) images_b_s = images_b_s.cuda(opts.gpu).detach() for o in range(n_obj): gt_s[o] = gt_s[o].cuda(opts.gpu).detach() gt_pos_s[o] = gt_pos_s[o].cuda(opts.gpu).detach() shutter_speed_s = shutter_speed_s.cuda(opts.gpu).detach() # test model s = 0 loss_pred_last_sum = 0 #last_gt = gt_s[:,opts.sequence_num - 1,:] last_gt = None z = None while (s < (opts.sequence_num - 2)): if (z_delete_mode): z = z_zero_test left_steps = opts.sequence_num - s - 2 # 현재(s+1)스텝으로부터 몇번까지 더 gt 가 있는지 images_prev = images_b_s[:, s, :, :, :] images_b = images_b_s[:, s + 1, :, :, :] images_post = images_b_s[:, s + 2, :, :, :] gt = [] gt_pos = [] for o in range(n_obj): gt.append(gt_s[o][:, s + 1, :]) gt_pos.append(gt_pos_s[o][:, s + 1, :]) next_gt = [] next_gt_pos = [] for o in range(n_obj): next_gt.append(gt_s[o][:, s + 2, :]) next_gt_pos.append(gt_pos_s[o][:, s + 2, :]) loss_pred_last_sum += model.loss_pred_last_vel shutter_speed = shutter_speed_s if throwing_mode and s != 0: given_vel = model.next_vel_pred given_pos = model.next_pos_pred else: given_vel = None given_pos = None model.test(images_prev, images_post, images_b, gt, gt_pos, next_gt, next_gt_pos, z, last_gt=last_gt, left_steps=left_steps, gt_pos_set=gt_pos_s, shutter_speed=shutter_speed, given_vel=given_vel, given_pos=given_pos) s += 1 total_dir_loss += model.loss_vel_dir[0] total_speed_loss += model.loss_vel_speed[0] total_vel_loss += model.loss_vel[0] total_pos_loss += model.loss_pos[0] total_dir_entropy += model.loss_dir_entropy[0] n += 1 if (model.pred_mode): z = model.z_next else: z = None #print("dir") #print(model.dir) total_dir_loss /= n total_speed_loss /= n total_vel_loss /= n total_pos_loss /= n total_dir_entropy /= n print('=============================') print("n: " + str(n)) if (vel_secret_mode): print( 'gen_loss: %04f, vel_loss %04f, vel_recons_loss %04f, gen_loss_gt_vel %04f, inverse_loss %04f, vel_dir_loss %04f, loss_content %04f, vel_pred_loss %04f, pos_loss %04f, dir_entropy_loss %04f, speed_loss %04f' % (model.loss_gen, total_vel_loss, model.loss_vel_recons, model.loss_gen_gt_vel, model.loss_inverse, total_dir_loss, model.loss_content, loss_pred_vel, total_pos_loss, total_dir_entropy, total_speed_loss)) else: if pred_mode: loss_pred_vel = model.loss_pred_vel[0] else: loss_pred_vel = -1 print( 'gen_loss: %04f, vel_loss %04f, vel_recons_loss %04f, gen_loss_gt_vel %04f, inverse_loss %04f, vel_dir_loss %04f, loss_content %04f, vel_pred_loss %04f, pos_loss %04f, dir_entropy_loss %04f, speed_loss %04f' % (model.loss_gen, total_vel_loss, model.loss_vel_recons, model.loss_gen_gt_vel, model.loss_inverse, total_dir_loss, model.loss_content, loss_pred_vel, total_pos_loss, total_dir_entropy, total_speed_loss)) print('=============================') if (model.gen_mode): if ep % 10 == 0 and (it + 1) % 60 == 0: saver.write_img(ep * len(train_loader) + (it + 1), model) if vel_secret_mode or encoder_tuning_mode: if (ep + 1) % 5 == 0: save_pos(model.input_gt_vel[0], model.vel_pred[0], model.vel_pred[0], model.next_vel_pred[0], model.next_vel_encoded[0], opts.visualize_root, "example%06d.png" % (ep * len(train_loader) + (it + 1)), gt_next_=model.next_vel_encoded[0], pred_next_=model.next_vel_pred[0], gt_pos_=model.next_pos_encoded[0], pred_pos_=model.next_pos_pred[0]) elif (model.gen_mode and model.pred_mode): if ep % 10 == 0 and (it + 1) % 60 == 0: save_pos(model.input_gt_vel[0], model.vel_pred[0], model.vel_pred[0], model.input_gt_vel[0], model.vel_pred[0], opts.visualize_root, "example%06d.png" % (ep * len(train_loader) + (it + 1)), gt_next_=model.next_gt_vel[0], pred_next_=model.next_vel_pred[0], gt_pos_=model.input_gt_pos[0], pred_pos_=model.pos_pred[0]) elif (model.pred_mode): if (it + 1) % 300 == 0: save_pos(model.input_gt_vel[0], model.vel_pred, model.vel_pred, model.input_gt_vel[0], model.vel_pred, opts.visualize_root, "example%06d.png" % (ep * len(train_loader) + (it + 1)), gt_next_=model.next_gt_vel, pred_next_=model.next_vel_pred, gt_pos_=model.input_gt_pos, pred_pos_=model.pos_pred) elif (model.gen_mode): if ep % 10 == 0 and (it + 1) % 60 == 0: save_pos(model.input_gt_vel[0], model.vel_pred[0], model.vel_pred[1], model.input_gt_vel[0], model.vel_pred[0], opts.visualize_root, "example%06d.png" % (ep * len(train_loader) + (it + 1)), gt_next_=model.input_gt_vel[0], pred_next_=model.vel_pred[0], gt_pos_=model.input_gt_pos[0], pred_pos_=model.pos_pred[0], gt_pos_2=model.input_gt_pos[0]) else: if ep % 10 == 0 and (it + 1) % 20 == 0: save_pos(model.input_gt_vel[0], model.vel_pred, model.vel_pred, model.input_gt_vel[0], model.vel_pred, opts.visualize_root, "example%06d.png" % (ep * len(train_loader) + (it + 1)), gt_next_=model.input_gt_vel[0], pred_next_=model.vel_pred, gt_pos_=model.input_gt_pos[0], pred_pos_=model.pos_pred, gt_pos_2=model.input_gt_pos[0]) #save_pos_for_next_gt(opts.visualize_root,"example%06d.png"%(ep*len(train_loader) + (it+1)), gt_next_ = model.next_gt, pred_next_ = model.next_gt_pred) total_it += 1 if total_it >= max_it: saver.write_img(-1, model) saver.write_model(-1, model) break # decay learning rate if opts.n_ep_decay > -1: model.update_lr() # Save network weights #if vel_secret_mode or encoder_tuning_mode: #saver.write_model(ep, total_it+1, model) if (ep + 1) % 5 == 0: saver.write_model(ep, total_it + 1, model) return
# Main module to train the model, load the data, # do gradient descent etc. followed by saving our model for later testing import os import numpy as np import torch from torchvision.transforms import * import torch.optim as optim from options import TrainOptions from model import Model from dataloader import DataSet # Get the Hyperparaeters opt = TrainOptions().parse() sample_dataset = DataSet(opt,"../dataset/IMFDB") sample_loader = torch.utils.data.DataLoader(sample_dataset,batch_size=opt.batch_size,num_workers=5,shuffle=False) # Check if gpu available or not device = torch.device("cuda" if (torch.cuda.is_available() and opt.use_gpu) else "cpu") opt.device = device # Load the model and send it to gpu model = Model(opt,15) model = model.to(device) if opt.use_gpu: model = torch.nn.DataParallel(model, device_ids=opt.gpus)
def main(): # parse options parser = TrainOptions() opts = parser.parse() # data loader print('--- load parameter ---') outer_iter = opts.outer_iter fade_iter = max(1.0, float(outer_iter / 2)) epochs = opts.epoch batchsize = opts.batchsize datasize = opts.datasize datarange = opts.datarange augementratio = opts.augementratio centercropratio = opts.centercropratio # model print('--- create model ---') tetGAN = TETGAN(gpu=(opts.gpu != 0)) if opts.gpu != 0: tetGAN.cuda() tetGAN.init_networks(weights_init) tetGAN.train() print('--- training ---') stylenames = os.listdir(opts.train_path) print('List of %d styles:' % (len(stylenames)), *stylenames, sep=' ') if opts.progressive == 1: # proressive training. From level1 64*64, to level2 128*128, to level3 256*256 # level 1 for i in range(outer_iter): jitter = min(1.0, i / fade_iter) fnames = load_trainset_batchfnames(opts.train_path, batchsize * 4, datarange, datasize * 2) for epoch in range(epochs): for fname in fnames: x, y_real, y = prepare_batch(fname, 1, jitter, centercropratio, augementratio, opts.gpu) losses = tetGAN.one_pass(x[0], None, y[0], None, y_real[0], None, 1, None) print('Level1, Iter[%d/%d], Epoch [%d/%d]' % (i + 1, outer_iter, epoch + 1, epochs)) print( 'Lrec: %.3f, Ldadv: %.3f, Ldesty: %.3f, Lsadv: %.3f, Lsty: %.3f' % (losses[0], losses[1], losses[2], losses[3], losses[4])) # level 2 for i in range(outer_iter): w = max(0.0, 1 - i / fade_iter) fnames = load_trainset_batchfnames(opts.train_path, batchsize * 2, datarange, datasize * 2) for epoch in range(epochs): for fname in fnames: x, y_real, y = prepare_batch(fname, 2, 1, centercropratio, augementratio, opts.gpu) losses = tetGAN.one_pass(x[0], x[1], y[0], y[1], y_real[0], y_real[1], 2, w) print('Level2, Iter[%d/%d], Epoch [%d/%d]' % (i + 1, outer_iter, epoch + 1, epochs)) print( 'Lrec: %.3f, Ldadv: %.3f, Ldesty: %.3f, Lsadv: %.3f, Lsty: %.3f' % (losses[0], losses[1], losses[2], losses[3], losses[4])) # level 3 for i in range(outer_iter): w = max(0.0, 1 - i / fade_iter) fnames = load_trainset_batchfnames(opts.train_path, batchsize, datarange, datasize) for epoch in range(epochs): for fname in fnames: x, y_real, y = prepare_batch(fname, 3, 1, centercropratio, augementratio, opts.gpu) losses = tetGAN.one_pass(x[0], x[1], y[0], y[1], y_real[0], y_real[1], 3, w) print('Level3, Iter[%d/%d], Epoch [%d/%d]' % (i + 1, outer_iter, epoch + 1, epochs)) print( 'Lrec: %.3f, Ldadv: %.3f, Ldesty: %.3f, Lsadv: %.3f, Lsty: %.3f' % (losses[0], losses[1], losses[2], losses[3], losses[4])) else: # directly train on level3 256*256 for i in range(outer_iter): fnames = load_trainset_batchfnames(opts.train_path, batchsize, datarange, datasize) for epoch in range(epochs): for fname in fnames: x, y_real, y = prepare_batch(fname, 3, 1, centercropratio, augementratio, opts.gpu) losses = tetGAN.one_pass(x[0], None, y[0], None, y_real[0], None, 3, 0) print('Iter[%d/%d], Epoch [%d/%d]' % (i + 1, outer_iter, epoch + 1, epochs)) print( 'Lrec: %.3f, Ldadv: %.3f, Ldesty: %.3f, Lsadv: %.3f, Lsty: %.3f' % (losses[0], losses[1], losses[2], losses[3], losses[4])) print('--- save ---') torch.save(tetGAN.state_dict(), opts.save_model_name)
def run_train(train_params=None, test_params=None): opt = TrainOptions().parse(train_params) testopt = TestOptions().parse(test_params) testopt.timestamp = opt.timestamp testopt.batch_size = 30 # model init model = SketchModel(opt) model.print_detail() writer = Writer(opt) # data load trainDataloader = load_data(opt, datasetType='train', permutation=opt.permutation, shuffle=opt.shuffle) validDataloader = load_data(opt, datasetType='valid') testDataloader = load_data(opt, datasetType='test') # train epoches # with torchsnooper.snoop(): ii = 0 min_test_avgloss = 100 min_test_avgloss_epoch = 0 for epoch in range(opt.epoch): for i, data in enumerate(trainDataloader): model.step(data) if ii % opt.plot_freq == 0: writer.plot_train_loss(model.loss, ii) if ii % opt.print_freq == 0: writer.print_train_loss(epoch, i, model.loss) ii += 1 model.update_learning_rate() if opt.plot_weights: writer.plot_model_wts(model, epoch) # test if epoch % opt.run_test_freq == 0: model.save_network('latest') loss_avg, P_metric, C_metric = run_eval(opt=testopt, loader=validDataloader, dataset='valid', write_result=False) writer.print_test_loss(epoch, 0, loss_avg) writer.plot_test_loss(loss_avg, epoch) writer.print_eval_metric(epoch, P_metric, C_metric) writer.plot_eval_metric(epoch, P_metric, C_metric) if loss_avg < min_test_avgloss: min_test_avgloss = loss_avg min_test_avgloss_epoch = epoch print( 'saving the model at the end of epoch {} with best avgLoss {}' .format(epoch, min_test_avgloss)) model.save_network('bestloss') testopt.which_epoch = 'latest' testopt.metric_way = 'wlen' loss_avg, P_metric, C_metric = run_eval(opt=testopt, loader=testDataloader, dataset='test', write_result=False) testopt.which_epoch = 'bestloss' testopt.metric_way = 'wlen' loss_avg_2, P_metric_2, C_metric_2 = run_eval(opt=testopt, loader=testDataloader, dataset='test', write_result=False) record_list = { 'p_metric': round(P_metric * 100, 2), 'c_metric': round(C_metric * 100, 2), 'loss_avg': round(loss_avg, 4), 'best_epoch': min_test_avgloss_epoch, 'p_metric_2': round(P_metric_2 * 100, 2), 'c_metric_2': round(C_metric_2 * 100, 2), 'loss_avg_2': round(loss_avg_2, 4), } writer.train_record(record_list=record_list) writer.close() return record_list, opt.timestamp
def main(): # parse options parser = TrainOptions() opts = parser.parse() # data loader print('\n--- load dataset ---') if opts.multi_modal: dataset = dataset_unpair_multi(opts) else: dataset = dataset_unpair(opts) train_loader = torch.utils.data.DataLoader(dataset, batch_size=opts.batch_size, shuffle=True, num_workers=opts.nThreads) # model print('\n--- load model ---') model = DRIT(opts) model.setgpu(opts.gpu) if opts.resume is None: model.initialize() ep0 = -1 total_it = 0 else: ep0, total_it = model.resume(opts.resume) model.set_scheduler(opts, last_ep=ep0) ep0 += 1 print('start the training at epoch %d' % (ep0)) # saver for display and output saver = Saver(opts) # train print('\n--- train ---') max_it = 500000 for ep in range(ep0, opts.n_ep): for it, (images_a, images_b) in enumerate(train_loader): if images_a.size(0) != opts.batch_size or images_b.size( 0) != opts.batch_size: continue # input data images_a = images_a.cuda(opts.gpu).detach() images_b = images_b.cuda(opts.gpu).detach() # update model if (it + 1) % opts.d_iter != 0 and it < len(train_loader) - 2: model.update_D_content(images_a, images_b) continue else: model.update_D(images_a, images_b) model.update_EG() # save to display file if not opts.no_display_img and not opts.multi_modal: saver.write_display(total_it, model) print('total_it: %d (ep %d, it %d), lr %08f' % (total_it, ep, it, model.gen_opt.param_groups[0]['lr'])) total_it += 1 if total_it >= max_it: # saver.write_img(-1, model) saver.write_model(-1, model) break # decay learning rate if opts.n_ep_decay > -1: model.update_lr() # save result image if not opts.multi_modal: saver.write_img(ep, model) # Save network weights saver.write_model(ep, total_it, model) return
def train_model(): opt = TrainOptions().parse(sub_dirs=['vis_multi','vis_cycle','vis_latest','train_vis_cycle']) out_f = open("%s/results.txt" % opt.expr_dir, 'w') copy_scripts_to_folder(opt.expr_dir) use_gpu = len(opt.gpu_ids) > 0 if opt.seed is not None: print ("using random seed:", opt.seed) random.seed(opt.seed) np.random.seed(opt.seed) torch.manual_seed(opt.seed) if use_gpu: torch.cuda.manual_seed_all(opt.seed) if opt.numpy_data: trainA, trainB, devA, devB, testA, testB = load_edges2shoes(opt.dataroot) train_dataset = UnalignedIterator(trainA, trainB, batch_size=opt.batchSize) print_log(out_f, '#training images = %d' % len(train_dataset)) vis_inf = False test_dataset = AlignedIterator(testA, testB, batch_size=100) print_log(out_f, '#test images = %d' % len(test_dataset)) dev_dataset = AlignedIterator(devA, devB, batch_size=100) print_log(out_f, '#dev images = %d' % len(dev_dataset)) dev_cycle = itertools.cycle(AlignedIterator(devA, devB, batch_size=25)) else: train_data_loader = DataLoader(opt, subset='train', unaligned=True, batchSize=opt.batchSize) test_data_loader = DataLoader(opt, subset='test', unaligned=False, batchSize=200) dev_data_loader = DataLoader(opt, subset='dev', unaligned=False, batchSize=200) dev_cycle_loader = DataLoader(opt, subset='dev', unaligned=False, batchSize=25) train_dataset = train_data_loader.load_data() dataset_size = len(train_data_loader) print_log(out_f, '#training images = %d' % dataset_size) vis_inf = False test_dataset = test_data_loader.load_data() print_log(out_f, '#test images = %d' % len(test_data_loader)) dev_dataset = dev_data_loader.load_data() print_log(out_f, '#dev images = %d' % len(dev_data_loader)) dev_cycle = itertools.cycle(dev_cycle_loader.load_data()) if opt.supervised: if opt.numpy_data: sup_size = int(len(trainA) * opt.sup_frac) sup_trainA = trainA[:sup_size] sup_trainB = trainB[:sup_size] sup_train_dataset = AlignedIterator(sup_trainA, sup_trainB, batch_size=opt.batchSize) else: sup_train_data_loader = DataLoader(opt, subset='train', unaligned=False, batchSize=opt.batchSize, fraction=opt.sup_frac) sup_train_dataset = sup_train_data_loader.load_data() sup_size = len(sup_train_data_loader) sup_train_dataset = itertools.cycle(sup_train_dataset) print_log(out_f, '#supervised images = %d' % sup_size) # create_model if opt.model == 'stoch_cycle_gan': model = StochCycleGAN(opt) elif opt.model == 'cycle_gan': model = StochCycleGAN(opt, ignore_noise=True) elif opt.model == 'aug_cycle_gan': model = AugmentedCycleGAN(opt) create_sub_dirs(opt, ['vis_inf']) vis_inf = True else: raise NotImplementedError('Specified model is not implemented.') print_log(out_f, "model [%s] was created" % (model.__class__.__name__)) # visualizer = Visualizer(opt) total_steps = 0 print_start_time = time.time() results = { 'best_dev_mse_A' : sys.float_info.max, 'best_test_mse_A' : sys.float_info.max, 'best_dev_bpp_B' : sys.float_info.max, 'best_test_bpp_B' : sys.float_info.max, } save_results(opt.expr_dir, results) history_mse_A = [] history_ubo_B = [] create_sub_dirs(opt, ['vis_pred_B']) for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1): epoch_start_time = time.time() epoch_iter = 0 for i, data in enumerate(train_dataset): real_A, real_B = Variable(data['A']), Variable(data['B']) if real_A.size(0) != real_B.size(0): continue prior_z_B = Variable(real_A.data.new(real_A.size(0), opt.nlatent, 1, 1).normal_(0, 1)) total_steps += opt.batchSize epoch_iter += opt.batchSize if use_gpu: real_A = real_A.cuda() real_B = real_B.cuda() prior_z_B = prior_z_B.cuda() if opt.monitor_gnorm: losses, visuals, gnorms = model.train_instance(real_A, real_B, prior_z_B) else: losses, visuals = model.train_instance(real_A, real_B, prior_z_B) # supervised training if opt.supervised: sup_data = sup_train_dataset.next() sup_real_A, sup_real_B = Variable(sup_data['A']), Variable(sup_data['B']) if use_gpu: sup_real_A, sup_real_B = sup_real_A.cuda(), sup_real_B.cuda() sup_losses = model.supervised_train_instance(sup_real_A, sup_real_B, prior_z_B) if total_steps % opt.display_freq == 0: # visualize current training batch visualize_cycle(opt, real_A, visuals, epoch, epoch_iter/opt.batchSize, train=True) # dev_data = dev_cycle.next() dev_data = next(dev_cycle) dev_real_A, dev_real_B = Variable(dev_data['A']), Variable(dev_data['B']) dev_prior_z_B = Variable(dev_real_A.data.new(dev_real_A.size(0), opt.nlatent, 1, 1).normal_(0, 1)) if use_gpu: dev_real_A = dev_real_A.cuda() dev_real_B = dev_real_B.cuda() dev_prior_z_B = dev_prior_z_B.cuda() dev_visuals = model.generate_cycle(dev_real_A, dev_real_B, dev_prior_z_B) visualize_cycle(opt, dev_real_A, dev_visuals, epoch, epoch_iter/opt.batchSize, train=False) # visualize generated B with different z_B visualize_multi(opt, dev_real_A, model, epoch, epoch_iter/opt.batchSize) if vis_inf: # visualize generated B with different z_B infered from real_B visualize_inference(opt, dev_real_A, dev_real_B, model, epoch, epoch_iter/opt.batchSize) if total_steps % opt.print_freq == 0: t = (time.time() - print_start_time) / opt.batchSize print_log(out_f, format_log(epoch, epoch_iter, losses, t)) if opt.supervised: print_log(out_f, format_log(epoch, epoch_iter, sup_losses, t, prefix=False)) if opt.monitor_gnorm: print_log(out_f, format_log(epoch, epoch_iter, gnorms, t, prefix=False)+"\n") print_start_time = time.time() if epoch % opt.save_epoch_freq == 0: print_log(out_f, 'saving the model at the end of epoch %d, iters %d' % (epoch, total_steps)) model.save('latest') ##################### # evaluate mappings ##################### if epoch % opt.eval_A_freq == 0: t = time.time() dev_mse_A = eval_mse_A(dev_dataset, model) test_mse_A = eval_mse_A(test_dataset, model) t = time.time() - t history_mse_A.append((dev_mse_A, test_mse_A)) np.save("%s/history_mse_A" % opt.expr_dir, history_mse_A) res_str_list = ["[%d] DEV_MSE_A: %.4f, TEST_MSE_A: %.4f, TIME: %.4f" % (epoch, dev_mse_A, test_mse_A, t)] if dev_mse_A < results['best_dev_mse_A']: with open("%s/best_mse_A.txt" % opt.expr_dir, 'w') as best_mse_A_f: best_mse_A_f.write(res_str_list[0]+'\n') best_mse_A_f.flush() results['best_dev_mse_A'] = dev_mse_A results['best_test_mse_A'] = test_mse_A model.save('best_A') save_results(opt.expr_dir, results) res_str_list += ["*** BEST DEV A ***"] res_str = "\n".join(["-"*60] + res_str_list + ["-"*60]) print_log(out_f, res_str) if epoch % opt.eval_B_freq == 0: t = time.time() if opt.model == 'cycle_gan': steps = 1 else: steps = 50 dev_ubo_B, dev_bpp_B, dev_kld_B = eval_ubo_B(dev_dataset, model, steps, True, 'pred_B_%d' % epoch, opt.vis_pred_B) test_ubo_B, test_bpp_B, test_kld_B = eval_ubo_B(test_dataset, model, steps, False, 'pred_B', opt.vis_pred_B) t = time.time() - t history_ubo_B.append((dev_ubo_B, dev_bpp_B, dev_kld_B, test_ubo_B, test_bpp_B, test_kld_B)) np.save("%s/history_ubo_B" % opt.expr_dir, history_ubo_B) res_str_list = ["[%d] DEV_BPP_B: %.4f, TEST_BPP_B: %.4f, TIME: %.4f" % (epoch, dev_bpp_B, test_bpp_B, t)] if dev_bpp_B < results['best_dev_bpp_B']: with open("%s/best_bpp_B.txt" % opt.expr_dir, 'w') as best_bpp_B_f: best_bpp_B_f.write(res_str_list[0]+'\n') best_bpp_B_f.flush() results['best_dev_bpp_B'] = dev_bpp_B results['best_test_bpp_B'] = test_bpp_B save_results(opt.expr_dir, results) model.save('best_B') res_str_list += ["*** BEST BPP B ***"] res_str = "\n".join(["-"*60] + res_str_list + ["-"*60]) print_log(out_f, res_str) print_log(out_f, 'End of epoch %d / %d \t Time Taken: %d sec' % (epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time)) if epoch > opt.niter: model.update_learning_rate() out_f.close()
def main(): parser = TrainOptions() opts = parser.parse() dataroot = opts.dataroot workers = opts.workers batch_size = opts.batch_size image_size = opts.image_size nc = opts.num_channels num_epochs = opts.num_epochs lr = opts.lr beta1 = opts.beta1 ngpu = opts.gpu checkpoint_path = opts.checkpoint_path device = torch.device("cuda:0" if ( torch.cuda.is_available() and ngpu > 0) else "cpu") nz = 128 ngf = 256 ndf = 256 def weights_init_normal(model): for param in model.parameters(): if (len(param.size()) == 2): torch.nn.init.xavier_normal(param) dataset = datasets.ImageFolder(root=dataroot, transform=transforms.Compose([ transforms.RandomRotation(0.5), transforms.RandomAffine(0.5), transforms.ColorJitter( 0, 0.1, 0.1, 0.1), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ])) dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=workers, pin_memory=True) netG = Generator(ngpu).to(device) if (device.type == 'cuda') and (ngpu > 1): netG = nn.DataParallel(netG, list(range(ngpu))) netG.apply(weights_init_normal) netD = Discriminator(ngpu).to(device) if (device.type == 'cuda') and (ngpu > 1): netD = nn.DataParallel(netD, list(range(ngpu))) netD.apply(weights_init_normal) criterion = nn.BCELoss() fixed_noise = torch.randn(32, nz, 1, 1, device=device) real_label = 1 fake_label = 0 optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999)) optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999)) # Training Loop print("Starting Training Loop...") for epoch in range(num_epochs): for i, data in enumerate(dataloader, 0): netD.zero_grad() real_cpu = data[0].to(device) b_size = real_cpu.size(0) label = torch.full((b_size, ), real_label, device=device) output = netD(real_cpu).view(-1) errD_real = criterion(output, label) errD_real.backward() D_x = output.mean().item() noise = torch.randn(b_size, nz, 1, 1, device=device) fake = netG(noise) label.fill_(fake_label) output = netD(fake.detach()).view(-1) errD_fake = criterion(output, label) errD_fake.backward() D_G_z1 = output.mean().item() errD = errD_real + errD_fake optimizerD.step() netG.zero_grad() label.fill_(real_label) # fake labels are real for generator cost output = netD(fake).view(-1) errG = criterion(output, label) errG.backward() D_G_z2 = output.mean().item() optimizerG.step() # Output training stats print( '[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f' % (epoch, num_epochs, i, len(dataloader), errD.item(), errG.item(), D_x, D_G_z1, D_G_z2)) torch.save(netG.state_dict(), os.path.join(checkpoint_path, 'generator.pkl'))