def main(): # parse options parser = TestOptions() opts = parser.parse() # data loader print('--- load data ---') style = load_image(opts.style_name) if opts.gpu != 0: style = to_var(style) if opts.c2s == 1: content = load_image(opts.content_name, opts.content_type) if opts.gpu != 0: content = to_var(content) # model print('--- load model ---') tetGAN = TETGAN() tetGAN.load_state_dict(torch.load(opts.model)) if opts.gpu != 0: tetGAN.cuda() tetGAN.eval() print('--- testing ---') if opts.c2s == 1: result = tetGAN(content, style) else: result = tetGAN.desty_forward(style) if opts.gpu != 0: result = to_data(result) print('--- save ---') # directory result_filename = os.path.join(opts.result_dir, opts.name) if not os.path.exists(opts.result_dir): os.mkdir(opts.result_dir) save_image(result[0], result_filename)
def main(): # parse options parser = FinetuneOptions() opts = parser.parse() # data loader print('--- load parameter ---') outer_iter = opts.outer_iter epochs = opts.epoch batchsize = opts.batchsize datasize = opts.datasize stylename = opts.style_name # model print('--- create model ---') tetGAN = TETGAN(gpu=(opts.gpu != 0)) if opts.gpu != 0: tetGAN.cuda() tetGAN.load_state_dict(torch.load(opts.load_model_name)) tetGAN.train() print('--- training ---') # supervised one shot learning if opts.supervise == 1: for i in range(outer_iter): fnames = load_oneshot_batchfnames(stylename, batchsize, datasize) for epoch in range(epochs): for fname in fnames: x, y_real, y = prepare_batch(fname, 3, 0, 0, 0, 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])) # unsupervised one shot learning else: for i in range(outer_iter): fnames = load_oneshot_batchfnames(stylename, batchsize, datasize) for epoch in range(epochs): for fname in fnames: # no ground truth x provided _, y_real, _ = prepare_batch(fname, 3, 0, 0, 0, opts.gpu) Lsrec = tetGAN.update_style_autoencoder(y_real[0]) for fname in fnames: # no ground truth x provided _, y_real, y = prepare_batch(fname, 3, 0, 0, 0, opts.gpu) with torch.no_grad(): x_auxiliary = tetGAN.desty_forward(y_real[0]) losses = tetGAN.one_pass(x_auxiliary, 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, Lsrec: %.3f' % (losses[0], losses[1], losses[2], losses[3], losses[4], Lsrec)) print('--- save ---') torch.save(tetGAN.state_dict(), opts.save_model_name)