def main(argv): (opts, args) = parser.parse_args(argv) torch.manual_seed(opts.seed) torch.cuda.manual_seed(opts.seed) if not os.path.exists(opts.output_folder): os.makedirs(opts.output_folder) # Load experiment setting config = get_config(opts.config) style_dim = config['gen']['style_dim'] opts.num_style = 1 if opts.style != '' else opts.num_style # Setup model and data loader trainer = MUNIT_Trainer(config) state_dict = torch.load(opts.checkpoint) trainer.gen_a.load_state_dict(state_dict['a']) trainer.gen_b.load_state_dict(state_dict['b']) trainer.cuda() trainer.eval() encode = trainer.gen_a.encode if opts.a2b else trainer.gen_b.encode # encode function style_encode = trainer.gen_b.encode if opts.a2b else trainer.gen_a.encode # encode function decode = trainer.gen_b.decode if opts.a2b else trainer.gen_a.decode # decode function transform = transforms.Compose([ transforms.Resize(config['new_size']), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) image = Variable(transform(Image.open( opts.input).convert('RGB')).unsqueeze(0).cuda(), volatile=True) style_image = Variable(transform(Image.open( opts.style).convert('RGB')).unsqueeze(0).cuda(), volatile=True) if opts.style != '' else None # Start testing style_rand = Variable(torch.randn(opts.num_style, style_dim, 1, 1).cuda(), volatile=True) content, _ = encode(image) if opts.style != '': _, style = style_encode(style_image) else: style = style_rand for j in range(opts.num_style): s = style[j].unsqueeze(0) outputs = decode(content, s) outputs = (outputs + 1) / 2. path = os.path.join(opts.output_folder, 'output{:03d}.jpg'.format(j)) vutils.save_image(outputs.data, path, padding=0, normalize=True) if not opts.output_only: # also save input images vutils.save_image(image.data, os.path.join(opts.output_folder, 'input.jpg'), padding=0, normalize=True)
def main(argv): (opts, args) = parser.parse_args(argv) torch.manual_seed(opts.seed) torch.cuda.manual_seed(opts.seed) if not os.path.exists(opts.output_folder): os.makedirs(opts.output_folder) # Load experiment setting config = get_config(opts.config) input_dim = config['new_size'] if opts.a2b else config['input_dim_b'] style_dim = config['gen']['style_dim'] # Setup model and data loader data_loader = get_data_loader_folder(opts.input_folder, 1, False, input_dim == 1, crop=False) trainer = MUNIT_Trainer(config) state_dict = torch.load(opts.checkpoint) trainer.gen_a.load_state_dict(state_dict['a']) trainer.gen_b.load_state_dict(state_dict['b']) trainer.cuda() trainer.eval() encode = trainer.gen_a.encode if opts.a2b else trainer.gen_b.encode # encode function decode = trainer.gen_b.decode if opts.a2b else trainer.gen_a.decode # decode function # Start testing style_fixed = Variable(torch.randn(opts.num_style, style_dim, 1, 1).cuda(), volatile=True) for i, images in enumerate(data_loader): images = Variable(images.cuda(), volatile=True) content, _ = encode(images) style = style_fixed if opts.synchronized else Variable( torch.randn(opts.num_style, style_dim, 1, 1).cuda(), volatile=True) for j in range(opts.num_style): s = style[j].unsqueeze(0) outputs = decode(content, s) outputs = (outputs + 1) / 2. path = os.path.join(opts.output_folder, 'input{:03d}_output{:03d}.jpg'.format(i, j)) vutils.save_image(outputs.data, path, padding=0, normalize=True) if not opts.output_only: # also save input images vutils.save_image(images.data, os.path.join(opts.output_folder, 'input{:03d}.jpg'.format(i)), padding=0, normalize=True)
opts = parser.parse_args() torch.manual_seed(opts.seed) torch.cuda.manual_seed(opts.seed) if not os.path.exists(opts.output_folder): os.makedirs(opts.output_folder) # Load experiment setting opts.num_style = 1 if opts.style != '' else opts.num_style if opts.trainer == 'MUNIT': from trainer import MUNIT_Trainer from utils import get_all_data_loaders, get_config config = get_config(opts.config) trainer = MUNIT_Trainer(config) elif opts.trainer == 'UNIT': from trainer import UNIT_Trainer from utils import get_config config = get_config(opts.config) trainer = UNIT_Trainer(config) elif opts.trainer == 'CDUNIT': from cd_trainer import CDUNIT_Trainer from cd_utils import get_all_data_loaders, get_config config = get_config(opts.config) trainer = CDUNIT_Trainer(config) elif opts.trainer == 'SECUNIT': from secunit_trainer import SECUNIT_Trainer from secunit_utils import get_all_data_loaders, get_config config = get_config(opts.config) trainer = SECUNIT_Trainer(config)
from trainer import MUNIT_Trainer import yaml from utils import get_config, get_loader, prepare_sub_folder, Timer, write_loss, write_2images parser = argparse.ArgumentParser() parser.add_argument('--config', type=str, default='config.yaml', help='Path to the config file.') parser.add_argument('--output_path', type=str, default='.', help="outputs path") opts = parser.parse_args() config = get_config(opts.config) max_iter = config['max_iter'] display_size = config['display_size'] trainer = MUNIT_Trainer(config) trainer.cuda() data_loader_a = get_loader(os.path.join(config['data_root'], 'b'), config['crop_image_height'], config['new_size'], config['batch_size'], 'afhq', 'train', 1) data_loader_b = get_loader(os.path.join(config['data_root'], 'a'), config['crop_image_height'], config['new_size'], config['batch_size'], 'afhq', 'train', 1) train_display_images_a = torch.stack([data_loader_a.dataset[i][0] for i in range(display_size)]).cuda() train_display_images_b = torch.stack([data_loader_b.dataset[i][0] for i in range(display_size)]).cuda() # Setup logger and output folders
config['snapshot_dir']=opts.snapshot_dir config['snapshot_save_iter']=opts.snapshot_save_iter config['sample_C']=opts.sample_C config['trim']=opts.trim config['sample_B']=opts.sample_B config['sample_D']=opts.sample_D config['sample_A']=opts.sample_A config['batch_size']=opts.batch_size config['transform_A']=opts.transform_A config['transform_B']=opts.transform_B config['transform_C']=opts.transform_C config['transform_D']=opts.transform_D config['weight_temp']=opts.weight_temp config['recon_x_cyc_w']=opts.recon_x_cyc_w # Setup model and data loader. trainer = MUNIT_Trainer(config, resume_epoch=opts.resume, snapshot_dir=opts.snapshot_dir) trainer.cuda() dataset_letters = eval(opts.dataset_letters) samples = list() dataset_probs = list() augmentation = list() for i in range(config['n_datasets']): samples.append(config['sample_' + dataset_letters[i]]) augmentation.append(config['transform_' + dataset_letters[i]]) train_loader_list, test_loader_list = get_all_data_loaders(config, config['n_datasets'], samples, augmentation, config['trim'],opts.dataset_letters) loader_sizes = list() for l in train_loader_list:
parser.add_argument('--trainer', type=str, default='MUNIT', help="MUNIT|UNIT") opts = parser.parse_args() torch.manual_seed(opts.seed) torch.cuda.manual_seed(opts.seed) # Load experiment setting config = get_config(opts.config) opts.num_style = 1 if opts.style != '' else opts.num_style # Setup model and data loader config['vgg_model_path'] = opts.output_path if opts.trainer == 'MUNIT': style_dim = config['gen']['style_dim'] trainer = MUNIT_Trainer(config) else: sys.exit("Only support MUNIT") try: state_dict = torch.load(opts.checkpoint) trainer.gen_a.load_state_dict(state_dict['a']) trainer.gen_b.load_state_dict(state_dict['b']) except: state_dict = pytorch03_to_pytorch04(torch.load(opts.checkpoint), opts.trainer) trainer.gen_a.load_state_dict(state_dict['a']) trainer.gen_b.load_state_dict(state_dict['b']) trainer.cuda() trainer.eval() encode = trainer.gen_a.encode if opts.a2b else trainer.gen_b.encode # encode function
def main(): parser = argparse.ArgumentParser() parser.add_argument('--config', type=str, default='configs/edges2handbags_folder.yaml', help='Path to the config file.') parser.add_argument('--output_path', type=str, default='.', help="outputs path") #resume option => [, default='730000'] parser.add_argument("--resume", default='150000', action="store_true") parser.add_argument('--trainer', type=str, default='MUNIT', help="MUNIT|UNIT") opts = parser.parse_args() cudnn.benchmark = True # Load experiment setting config = get_config(opts.config) max_iter = config['max_iter'] display_size = config['display_size'] config['vgg_model_path'] = opts.output_path # Setup model and data loader if opts.trainer == 'MUNIT': trainer = MUNIT_Trainer(config) elif opts.trainer == 'UNIT': trainer = UNIT_Trainer(config) else: sys.exit("Only support MUNIT|UNIT") trainer.cuda() train_loader_a, train_loader_b, test_loader_a, test_loader_b = get_all_data_loaders( config) train_display_images_a = torch.stack( [train_loader_a.dataset[i] for i in range(display_size)]).cuda() train_display_images_b = torch.stack( [train_loader_b.dataset[i] for i in range(display_size)]).cuda() test_display_images_a = torch.stack( [test_loader_a.dataset[i] for i in range(display_size)]).cuda() test_display_images_b = torch.stack( [test_loader_b.dataset[i] for i in range(display_size)]).cuda() # Setup logger and output folders model_name = os.path.splitext(os.path.basename(opts.config))[0] train_writer = tensorboardX.SummaryWriter( os.path.join(opts.output_path + "/logs", model_name)) output_directory = os.path.join(opts.output_path + "/outputs", model_name) checkpoint_directory, image_directory = prepare_sub_folder( output_directory) shutil.copy(opts.config, os.path.join( output_directory, 'config.yaml')) # copy config file to output folder # Start training iterations = trainer.resume(checkpoint_directory, hyperparameters=config) if opts.resume else 0 while True: for it, (images_a, images_b) in enumerate(zip(train_loader_a, train_loader_b)): trainer.update_learning_rate() images_a, images_b = images_a.cuda().detach(), images_b.cuda( ).detach() with Timer("Elapsed time in update: %f"): # Main training code trainer.dis_update(images_a, images_b, config) trainer.gen_update(images_a, images_b, config) torch.cuda.synchronize() # Dump training stats in log file if (iterations + 1) % config['log_iter'] == 0: print("Iteration: %08d/%08d" % (iterations + 1, max_iter)) write_loss(iterations, trainer, train_writer) # Write images if (iterations + 1) % config['image_save_iter'] == 0: with torch.no_grad(): test_image_outputs = trainer.sample( test_display_images_a, test_display_images_b) train_image_outputs = trainer.sample( train_display_images_a, train_display_images_b) write_2images(test_image_outputs, display_size, image_directory, 'test_%08d' % (iterations + 1)) write_2images(train_image_outputs, display_size, image_directory, 'train_%08d' % (iterations + 1)) # HTML write_html(output_directory + "/index.html", iterations + 1, config['image_save_iter'], 'images') if (iterations + 1) % config['image_display_iter'] == 0: with torch.no_grad(): image_outputs = trainer.sample(train_display_images_a, train_display_images_b) write_2images(image_outputs, display_size, image_directory, 'train_current') # Save network weights if (iterations + 1) % config['snapshot_save_iter'] == 0: trainer.save(checkpoint_directory, iterations) iterations += 1 if iterations >= max_iter: sys.exit('Finish training')
torch.manual_seed(opts.seed) torch.cuda.manual_seed(opts.seed) # if not os.path.exists(opts.output_folder): # os.makedirs(opts.output_folder) # Load experiment setting config = get_config(opts.config) config['gen']['no_style_enc'] = opts.no_style_enc # Setup model and data loader config['vgg_w'] = 0 if opts.trainer == 'MUNIT': config['gen']['style_dim'] = config['gen']['z_num'] trainer = MUNIT_Trainer(config) elif opts.trainer == 'UNIT': trainer = UNIT_Trainer(config) else: sys.exit("Only support MUNIT|UNIT") ############## configure checkpoint from output_folder checkpoint_path = find_latest_model_file(os.path.join(opts.output_folder, 'checkpoints'), opts.checkpoint, keyword='gen') try: state_dict = torch.load(checkpoint_path) trainer.gen_a.load_state_dict(state_dict['a']) trainer.gen_b.load_state_dict(state_dict['b'])
opts = parser.parse_args() if comet_exp is not None: comet_exp.log_asset(file_data=opts.config, file_name="config.yaml") comet_exp.log_parameter("git_hash", opts.git_hash) cudnn.benchmark = True # Load experiment setting config = get_config(opts.config) max_iter = config["max_iter"] display_size = config["display_size"] config["vgg_model_path"] = opts.output_path # Setup model and data loader if opts.trainer == "MUNIT": trainer = MUNIT_Trainer(config) elif opts.trainer == "UNIT": trainer = UNIT_Trainer(config) else: sys.exit("Only support MUNIT|UNIT") trainer.cuda() train_loader_a, train_loader_b, test_loader_a, test_loader_b = get_all_data_loaders( config) if config["semantic_w"] > 0: train_loader_a_w_mask = get_data_loader_mask_and_im( config["data_list_train_a"], config["data_list_train_a_seg"], config["batch_size"], True,
opts.config = 'configs/handwriting_online.yaml' opts.check_files = True print("Running on Galois, config {}, check_files {}".format(opts.config, opts.check_files)) cudnn.benchmark = True # Load experiment setting config = get_config(opts.config) print(config) max_iter = config['max_iter'] display_size = config['display_size'] config['vgg_model_path'] = opts.output_path # Setup model and data loader if opts.trainer == 'MUNIT': trainer = MUNIT_Trainer(config) elif opts.trainer == 'UNIT': trainer = UNIT_Trainer(config) else: sys.exit("Only support MUNIT|UNIT") #trainer.cuda() #train_loader_a, train_loader_b, test_loader_a, test_loader_b, folders = get_all_data_loaders(config) (train_loader_a, tr_a), (train_loader_b, tr_b), (test_loader_a, test_a), (test_loader_b, test_b), folders = get_all_data_loaders_better(config) if opts.check_files and False: print("Checking files...") for folder in folders: print(folder) utils.check_files(folder)
def main(): from utils import get_all_data_loaders, prepare_sub_folder, write_html, write_loss, get_config, write_2images, Timer import argparse from torch.autograd import Variable from trainer import MUNIT_Trainer, UNIT_Trainer import torch.backends.cudnn as cudnn import torch # try: # from itertools import izip as zip # except ImportError: # will be 3.x series # pass import os import sys import tensorboardX import shutil os.environ["CUDA_VISIBLE_DEVICES"] = str(0) parser = argparse.ArgumentParser() parser.add_argument('--config', type=str, default='configs/edges2handbags_folder.yaml', help='Path to the config file.') parser.add_argument('--output_path', type=str, default='.', help="outputs path") parser.add_argument("--resume", action="store_true") parser.add_argument('--trainer', type=str, default='MUNIT', help="MUNIT|UNIT") opts = parser.parse_args() cudnn.benchmark = True ''' Note: https://www.pytorchtutorial.com/when-should-we-set-cudnn-benchmark-to-true/ 大部分情况下,设置这个 flag 可以让内置的 cuDNN 的 auto-tuner 自动寻找最适合当前配置的高效算法,来达到优化运行效率的问题 1. 如果网络的输入数据维度或类型上变化不大,设置 torch.backends.cudnn.benchmark = true 可以增加运行效率; 2. 如果网络的输入数据在每次 iteration 都变化的话,会导致 cnDNN 每次都会去寻找一遍最优配置,这样反而会降低运行效率。 ''' # Load experiment setting config = get_config(opts.config) max_iter = config['max_iter'] display_size = config['display_size'] config['vgg_model_path'] = opts.output_path # Setup model and data loader if opts.trainer == 'MUNIT': trainer = MUNIT_Trainer(config) elif opts.trainer == 'UNIT': trainer = UNIT_Trainer(config) else: sys.exit("Only support MUNIT|UNIT") trainer.cuda() train_loader_a, train_loader_b, test_loader_a, test_loader_b = get_all_data_loaders( config) train_display_images_a = torch.stack( [train_loader_a.dataset[i] for i in range(display_size)]).cuda() train_display_images_b = torch.stack( [train_loader_b.dataset[i] for i in range(display_size)]).cuda() test_display_images_a = torch.stack( [test_loader_a.dataset[i] for i in range(display_size)]).cuda() test_display_images_b = torch.stack( [test_loader_b.dataset[i] for i in range(display_size)]).cuda() # Setup logger and output folders model_name = os.path.splitext(os.path.basename(opts.config))[0] train_writer = tensorboardX.SummaryWriter( os.path.join(opts.output_path + "/logs", model_name)) output_directory = os.path.join(opts.output_path + "/outputs", model_name) checkpoint_directory, image_directory = prepare_sub_folder( output_directory) shutil.copy(opts.config, os.path.join( output_directory, 'config.yaml')) # copy config file to output folder # Start training iterations = trainer.resume(checkpoint_directory, hyperparameters=config) if opts.resume else 0 while True: for it, (images_a, images_b) in enumerate(zip(train_loader_a, train_loader_b)): trainer.update_learning_rate() images_a, images_b = images_a.cuda().detach(), images_b.cuda( ).detach() with Timer("Elapsed time in update: %f"): # Main training code trainer.dis_update(images_a, images_b, config) trainer.gen_update(images_a, images_b, config) torch.cuda.synchronize() # Dump training stats in log file if (iterations + 1) % config['log_iter'] == 0: print("Iteration: %08d/%08d" % (iterations + 1, max_iter)) write_loss(iterations, trainer, train_writer) # Write images if (iterations + 1) % config['image_save_iter'] == 0: with torch.no_grad(): test_image_outputs = trainer.sample( test_display_images_a, test_display_images_b) train_image_outputs = trainer.sample( train_display_images_a, train_display_images_b) write_2images(test_image_outputs, display_size, image_directory, 'test_%08d' % (iterations + 1)) write_2images(train_image_outputs, display_size, image_directory, 'train_%08d' % (iterations + 1)) # HTML write_html(output_directory + "/index.html", iterations + 1, config['image_save_iter'], 'images') if (iterations + 1) % config['image_display_iter'] == 0: with torch.no_grad(): image_outputs = trainer.sample(train_display_images_a, train_display_images_b) write_2images(image_outputs, display_size, image_directory, 'train_current') # Save network weights if (iterations + 1) % config['snapshot_save_iter'] == 0: trainer.save(checkpoint_directory, iterations) iterations += 1 if iterations >= max_iter: sys.exit('Finish training')
def setup(opts): generator_checkpoint_path = opts['generator_checkpoint'] # generator_checkpoint_path = './checkpoints/ffhq2ladiescrop.pt' # Load experiment settings config = { 'image_save_iter': 10000, 'image_display_iter': 100, 'display_size': 16, 'snapshot_save_iter': 10000, 'log_iter': 100, 'max_iter': 1000000, 'batch_size': 1, 'weight_decay': 0.0001, 'beta1': 0.5, 'beta2': 0.999, 'init': 'kaiming', 'lr': 0.0001, 'lr_policy': 'step', 'step_size': 100000, 'gamma': 0.5, 'gan_w': 1, 'recon_x_w': 10, 'recon_s_w': 1, 'recon_c_w': 1, 'recon_x_cyc_w': 10, 'vgg_w': 0, 'gen': { 'dim': 64, 'mlp_dim': 256, 'style_dim': 8, 'activ': 'relu', 'n_downsample': 2, 'n_res': 4, 'pad_type': 'reflect' }, 'dis': { 'dim': 64, 'norm': 'none', 'activ': 'lrelu', 'n_layer': 4, 'gan_type': 'lsgan', 'num_scales': 3, 'pad_type': 'reflect' }, 'input_dim_a': 3, 'input_dim_b': 3, 'num_workers': 8, 'new_size': 1024, 'crop_image_height': 400, 'crop_image_width': 400, 'data_root': './datasets/ffhq2ladies/' } # Setup model and data loader trainer = MUNIT_Trainer(config) state_dict = torch.load(generator_checkpoint_path) trainer.gen_a.load_state_dict(state_dict['a']) trainer.gen_b.load_state_dict(state_dict['b']) return {'model': trainer, 'config': config}
list_images.append(image) list_classes_to_take.remove(label) # print(list_classes_to_take) return torch.stack(list_images).cuda() # train_display_images_a = torch.stack([loader.dataset[i][0]]).cuda() # Load experiment setting config = get_config(opts.config) max_iter = config['max_iter'] display_size = config['display_size'] config['vgg_model_path'] = opts.output_path # Setup model and data loader if opts.trainer == 'MUNIT': trainer = MUNIT_Trainer(config) # elif opts.trainer == 'UNIT': # trainer = UNIT_Trainer(config) else: sys.exit("Only support MUNIT|UNIT") trainer.cuda() train_loader_a, train_loader_a_limited, train_loader_b, test_loader_a, test_loader_b = get_all_data_loaders( config) # train_display_images_a_temp = torch.stack([train_loader_a.dataset[i][0] for i in range(display_size)]).cuda() # train_display_images_b = torch.stack([train_loader_b.dataset[i][0] for i in range(display_size)]).cuda() # test_display_images_a = torch.stack([test_loader_a.dataset[i][0] for i in range(display_size)]).cuda() # test_display_images_b = torch.stack([test_loader_b.dataset[i][0] for i in range(display_size)]).cuda() train_display_images_a = get_display_images(train_loader_a) train_display_images_b = get_display_images(train_loader_b) test_display_images_a = get_display_images(test_loader_a)
parser.add_argument('--output_path', type=str, default='.', help="outputs path") parser.add_argument("--resume", action="store_true") parser.add_argument('--trainer', type=str, default='MUNIT', help="MUNIT|UNIT") opts = parser.parse_args() cudnn.benchmark = True # Load experiment setting config = get_config(opts.config) max_iter = config['max_iter'] display_size = config['display_size'] config['vgg_model_path'] = opts.output_path # Setup model and data loader if opts.trainer == 'MUNIT': trainer = MUNIT_Trainer(config) elif opts.trainer == 'UNIT': trainer = UNIT_Trainer(config) else: sys.exit("Only support MUNIT|UNIT") trainer.cuda() train_loader_a, train_loader_b, test_loader_a, test_loader_b = get_all_data_loaders(config) train_display_images_a = Variable(torch.stack([train_loader_a.dataset[i] for i in range(display_size)]).cuda(), volatile=True) train_display_images_b = Variable(torch.stack([train_loader_b.dataset[i] for i in range(display_size)]).cuda(), volatile=True) test_display_images_a = Variable(torch.stack([test_loader_a.dataset[i] for i in range(display_size)]).cuda(), volatile=True) test_display_images_b = Variable(torch.stack([test_loader_b.dataset[i] for i in range(display_size)]).cuda(), volatile=True) # Setup logger and output folders model_name = os.path.splitext(os.path.basename(opts.config))[0] train_writer = tensorboardX.SummaryWriter(os.path.join(opts.output_path + "/logs", model_name)) output_directory = os.path.join(opts.output_path + "/outputs", model_name)
config['no_rec_s'] = opts.no_rec_s if opts.vgg_w != -1: config['vgg_w'] = opts.vgg_w if opts.cyc_rec_weight != -1: config['recon_x_cyc_w'] = opts.cyc_rec_weight if opts.ne_weight != -1: config['loss_eg_weight'] = opts.ne_weight if opts.rec_c_weight != -1: config['recon_c_w'] = opts.rec_c_weight # Setup model and data loader if opts.trainer == 'UNIT': trainer = UNIT_Trainer(config) elif opts.trainer == 'MUNIT': trainer = MUNIT_Trainer(config) trainer.cuda() Dataset = choose_dataset(config['dataset_name']) dataset = Dataset(config['data_root'], config, split='train') dataloader = torch.utils.data.DataLoader(dataset, batch_size=config['batch_size'], drop_last=True, shuffle=True, num_workers=int( config['num_workers'])) # Setup logger and output folders model_name = os.path.splitext(os.path.basename(opts.config))[0] + '_vgg_%s_%s' \
help="outputs path") parser.add_argument('--load', type=int, default=400) parser.add_argument('--snapshot_dir', type=str, default='.') opts = parser.parse_args() cudnn.benchmark = True # Load experiment setting. config = get_config(opts.config) display_size = config['display_size'] config['vgg_model_path'] = opts.output_path # Setup model and data loader. if config['trainer'] == 'MUNIT': trainer = MUNIT_Trainer(config, resume_epoch=opts.load, snapshot_dir=opts.snapshot_dir) elif config['trainer'] == 'UNIT': trainer = UNIT_Trainer(config, resume_epoch=opts.load, snapshot_dir=opts.snapshot_dir) else: sys.exit("Only support MUNIT|UNIT.") os.exit() trainer.cuda() dataset_letters = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I'] samples = list() dataset_probs = list() augmentation = list()
torch.manual_seed(opts.seed) torch.cuda.manual_seed(opts.seed) if not os.path.exists(opts.output_folder): os.makedirs(opts.output_folder) # Load experiment setting config = get_config(opts.config) opts.num_style = 1 if opts.style != '' else opts.num_style # Setup model and data loader if opts.trainer == 'MUNIT': style_dim = config['gen']['style_dim'] trainer = MUNIT_Trainer(config) elif opts.trainer == 'UNIT': trainer = UNIT_Trainer(config) elif opts.trainer == 'AttnMUNIT': trainer = AttnMUNIT_Trainer(config,opts.discriminator,opts.attention,opts.concat_type) else: sys.exit("Only support AttnMUNIT|MUNIT|UNIT") # elif opts.trainer == 'AdaINAttnMUNIT': # trainer = AdaINAttnMUNIT_Trainer(config) try: state_dict = torch.load(opts.checkpoint) trainer.gen_a.load_state_dict(state_dict['a']) trainer.gen_b.load_state_dict(state_dict['b']) except: state_dict = pytorch03_to_pytorch04(torch.load(opts.checkpoint), opts.trainer)
torch.manual_seed(opts.seed) torch.cuda.manual_seed(opts.seed) if not os.path.exists(opts.output_folder): os.makedirs(opts.output_folder) device = 'cuda:%d'%opts.gpu_id # Load experiment setting config = get_config(opts.config) opts.num_style = 1 if opts.style != '' else opts.num_style # Setup model and data loader config['vgg_model_path'] = opts.output_path if opts.trainer == 'MUNIT': style_dim = config['gen']['style_dim'] trainer = MUNIT_Trainer(config, device) elif opts.trainer == 'UNIT': trainer = UNIT_Trainer(config, device) else: sys.exit("Only support MUNIT|UNIT") try: state_dict = torch.load(opts.checkpoint) trainer.gen_a.load_state_dict(state_dict['a']) trainer.gen_b.load_state_dict(state_dict['b']) except: state_dict = pytorch03_to_pytorch04(torch.load(opts.checkpoint), opts.trainer) trainer.gen_a.load_state_dict(state_dict['a']) trainer.gen_b.load_state_dict(state_dict['b']) trainer.cuda(device)
parser.add_argument('--config', type=str, default='configs/CXR_lungs', help='Path to the config file.') parser.add_argument('--output_path', type=str, default='.', help="Outputs path.") parser.add_argument('--resume', type=int, default=-1) parser.add_argument('--snapshot_dir', type=str, default='.') opts = parser.parse_args() cudnn.benchmark = True # Load experiment setting. config = get_config(opts.config) display_size = config['display_size'] config['vgg_model_path'] = opts.output_path # Setup model and data loader. if config['trainer'] == 'MUNIT': trainer = MUNIT_Trainer(config, resume_epoch=opts.resume, snapshot_dir=opts.snapshot_dir) elif config['trainer'] == 'UNIT': trainer = UNIT_Trainer(config, resume_epoch=opts.resume, snapshot_dir=opts.snapshot_dir) else: sys.exit("Only support MUNIT|UNIT.") os.exit() trainer.cuda() dataset_letters = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I'] samples = list() dataset_probs = list() augmentation = list() for i in range(config['n_datasets']): samples.append(config['sample_' + dataset_letters[i]]) dataset_probs.append(config['prob_' + dataset_letters[i]])
torch.manual_seed(opts.seed) torch.cuda.manual_seed(opts.seed) # Load experiment setting config = get_config(opts.config) input_dim = config['input_dim_a'] if opts.a2b else config['input_dim_b'] # Setup model and data loader image_names = ImageFolder(opts.input_folder, transform=None, return_paths=True) data_loader = get_data_loader_folder(opts.input_folder, 1, False, new_size=config['new_size_a'], crop=False) config['vgg_model_path'] = opts.output_path if opts.trainer == 'MUNIT': style_dim = config['gen']['style_dim'] trainer = MUNIT_Trainer(config) elif opts.trainer == 'UNIT': trainer = UNIT_Trainer(config) else: sys.exit("Only support MUNIT|UNIT") state_dict = torch.load(opts.checkpoint) trainer.gen_a.load_state_dict(state_dict['a']) trainer.gen_b.load_state_dict(state_dict['b']) trainer.cuda() trainer.eval() encode = trainer.gen_a.encode if opts.a2b else trainer.gen_b.encode # encode function decode = trainer.gen_b.decode if opts.a2b else trainer.gen_a.decode # decode function if opts.trainer == 'MUNIT':
opts = parser.parse_args() if comet_exp is not None: comet_exp.log_asset(file_data=opts.config, file_name="config.yaml") comet_exp.log_parameter("git_hash", opts.git_hash) cudnn.benchmark = True # Load experiment setting config = get_config(opts.config) max_iter = config["max_iter"] display_size = config["display_size"] config["vgg_model_path"] = opts.output_path # Setup model and data loader if opts.trainer == "MUNIT": trainer = MUNIT_Trainer(config) elif opts.trainer == "UNIT": trainer = UNIT_Trainer(config) else: sys.exit("Only support MUNIT|UNIT") trainer.cuda() print(config) train_loader_a, train_loader_b, test_loader_a, test_loader_b = get_all_data_loaders( config) #If if config["semantic_w"] > 0: train_loader_a_w_mask = get_data_loader_mask_and_im( config["data_list_train_a"],
def main(argv): (opts, args) = parser.parse_args(argv) cudnn.benchmark = True model_name = os.path.splitext(os.path.basename(opts.config))[0] # Load experiment setting config = get_config(opts.config) max_iter = config['max_iter'] display_size = config['display_size'] # Setup model and data loader trainer = MUNIT_Trainer(config) trainer.cuda() train_loader_a, train_loader_b, test_loader_a, test_loader_b = get_all_data_loaders( config) test_display_images_a = Variable(torch.stack( [test_loader_a.dataset[i] for i in range(display_size)]).cuda(), volatile=True) test_display_images_b = Variable(torch.stack( [test_loader_b.dataset[i] for i in range(display_size)]).cuda(), volatile=True) train_display_images_a = Variable(torch.stack( [train_loader_a.dataset[i] for i in range(display_size)]).cuda(), volatile=True) train_display_images_b = Variable(torch.stack( [train_loader_b.dataset[i] for i in range(display_size)]).cuda(), volatile=True) # Setup logger and output folders train_writer = tensorboard.SummaryWriter(os.path.join( opts.log, model_name)) output_directory = os.path.join(opts.outputs, model_name) checkpoint_directory, image_directory = prepare_sub_folder( output_directory) shutil.copy(opts.config, os.path.join( output_directory, 'config.yaml')) # copy config file to output folder # Start training iterations = trainer.resume(checkpoint_directory) if opts.resume else 0 while True: for it, (images_a, images_b) in enumerate(izip(train_loader_a, train_loader_b)): trainer.update_learning_rate() images_a, images_b = Variable(images_a.cuda()), Variable( images_b.cuda()) # Main training code trainer.dis_update(images_a, images_b, config) trainer.gen_update(images_a, images_b, config) # Dump training stats in log file if (iterations + 1) % config['log_iter'] == 0: print("Iteration: %08d/%08d" % (iterations + 1, max_iter)) write_loss(iterations, trainer, train_writer) # Write images if (iterations + 1) % config['image_save_iter'] == 0: # Test set images image_outputs = trainer.sample(test_display_images_a, test_display_images_b) write_images( image_outputs, display_size, '%s/gen_test%08d.jpg' % (image_directory, iterations + 1)) # Train set images image_outputs = trainer.sample(train_display_images_a, train_display_images_b) write_images( image_outputs, display_size, '%s/gen_train%08d.jpg' % (image_directory, iterations + 1)) # HTML write_html(output_directory + "/index.html", iterations + 1, config['image_save_iter'], 'images') if (iterations + 1) % config['image_save_iter'] == 0: image_outputs = trainer.sample(test_display_images_a, test_display_images_b) write_images(image_outputs, display_size, '%s/gen.jpg' % image_directory) # Save network weights if (iterations + 1) % config['snapshot_save_iter'] == 0: trainer.save(checkpoint_directory, iterations) iterations += 1 if iterations >= max_iter: return
for param in inception.parameters(): param.requires_grad = False inception_up = nn.Upsample(size=(299, 299), mode='bilinear') # Setup model and data loader image_names = ImageFolder(opts.input_folder, transform=None, return_paths=True) data_loader = get_data_loader_folder(opts.input_folder, 1, False, new_size=config['new_size_a'], crop=False) config['vgg_model_path'] = opts.output_path if opts.trainer == 'MUNIT': style_dim = config['gen']['style_dim'] trainer = MUNIT_Trainer(config) elif opts.trainer == 'UNIT': trainer = UNIT_Trainer(config) else: sys.exit("Only support MUNIT|UNIT") try: state_dict = torch.load(opts.checkpoint) trainer.gen_a.load_state_dict(state_dict['a']) trainer.gen_b.load_state_dict(state_dict['b']) except: state_dict = pytorch03_to_pytorch04(torch.load(opts.checkpoint), opts.trainer) trainer.gen_a.load_state_dict(state_dict['a']) trainer.gen_b.load_state_dict(state_dict['b'])
type=str, default='MUNIT', help="MUNIT|UNIT") opts = parser.parse_args() cudnn.benchmark = True # Load experiment setting config = get_config(opts.config) max_iter = config['max_iter'] display_size = config['display_size'] config['vgg_model_path'] = opts.output_path # Setup model and data loader if opts.trainer == 'MUNIT': trainer = MUNIT_Trainer(config) elif opts.trainer == 'UNIT': trainer = UNIT_Trainer(config) else: sys.exit("Only support MUNIT|UNIT") trainer.cuda() train_loader_a, train_loader_b, test_loader_a, test_loader_b = get_all_data_loaders( config) train_display_images_a = torch.stack( [train_loader_a.dataset[i] for i in range(display_size)]).cuda() train_display_images_b = torch.stack( [train_loader_b.dataset[i] for i in range(display_size)]).cuda() test_display_images_a = torch.stack( [test_loader_a.dataset[i] for i in range(display_size)]).cuda() test_display_images_b = torch.stack( [test_loader_b.dataset[i] for i in range(display_size)]).cuda()
def main(argv): (opts, args) = parser.parse_args(argv) cudnn.benchmark = True # Load experiment setting config = get_config(opts.config) max_iter = config['max_iter'] # Setup logger and output folders output_subfolders = prepare_logging_folders(config['output_root'], config['experiment_name']) logger = create_logger( os.path.join(output_subfolders['logs'], 'train_log.log')) shutil.copy(opts.config, os.path.join( output_subfolders['logs'], 'config.yaml')) # copy config file to output folder tb_logger = tensorboard_logger.Logger(output_subfolders['logs']) logger.info('============ Initialized logger ============') logger.info('Config File: {}'.format(opts.config)) # Setup model and data loader trainer = MUNIT_Trainer(config, opts) trainer.cuda() loaders = get_all_data_loaders(config) val_display_images = next(iter(loaders['val'])) logger.info('Test images: {}'.format(val_display_images['A_paths'])) # Start training iterations = trainer.resume(opts.model_path, hyperparameters=config) if opts.resume else 0 while True: for it, images in enumerate(loaders['train']): trainer.update_learning_rate() images_a = images['A'] images_b = images['B'] images_a, images_b = Variable(images_a.cuda()), Variable( images_b.cuda()) # Main training code trainer.dis_update(images_a, images_b, config) trainer.gen_update(images_a, images_b, config) # Dump training stats in log file if (iterations + 1) % config['log_iter'] == 0: for tag, value in trainer.loss.items(): tb_logger.scalar_summary(tag, value, iterations) val_output_imgs = trainer.sample( Variable(val_display_images['A'].cuda()), Variable(val_display_images['B'].cuda())) tb_imgs = [] for imgs in val_output_imgs.values(): tb_imgs.append(torch.cat(torch.unbind(imgs, 0), dim=2)) tb_logger.image_summary(list(val_output_imgs.keys()), tb_imgs, iterations) if (iterations + 1) % config['print_iter'] == 0: logger.info( "Iteration: {:08}/{:08} Discriminator Loss: {:.4f} Generator Loss: {:.4f}" .format(iterations + 1, max_iter, trainer.loss['D/total'], trainer.loss['G/total'])) # Write images # if (iterations + 1) % config['image_save_iter'] == 0: # val_output_imgs = trainer.sample( # Variable(val_display_images['A'].cuda()), # Variable(val_display_images['B'].cuda())) # # for key, imgs in val_output_imgs.items(): # key = key.replace('/', '_') # write_images(imgs, config['display_size'], '{}/{}_{:08}.jpg'.format(output_subfolders['images'], key, iterations+1)) # # logger.info('Saved images to: {}'.format(output_subfolders['images'])) # Save network weights if (iterations + 1) % config['snapshot_save_iter'] == 0: trainer.save(output_subfolders['models'], iterations) iterations += 1 if iterations >= max_iter: return