def main(): parser = argparse.ArgumentParser(description='chainer implementation of pix2pix') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--img', '-i', help='Input image') parser.add_argument('--out', '-o', default='result_dehighlight', help='Directory to output the result') args = parser.parse_args() # Set up a neural network to train enc = Encoder(in_ch=3) dec = Decoder(out_ch=1) chainer.serializers.load_npz(ENC_W, enc) chainer.serializers.load_npz(DEC_W, dec) if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current enc.to_gpu() # Copy the model to the GPU dec.to_gpu() inimg = loadimg(args.img) ch, h, w = inimg.shape # add paddings so that input array has the size of mutiples of 256. in_ary = np.zeros((ch,math.ceil(h/256)*256, math.ceil(w/256)*256), dtype="f") in_ary[:,0:h,0:w] = inimg x_in = in_ary[np.newaxis,:] # to fit into the minibatch shape print(x_in.shape) # x_in as an input image x_in = chainer.Variable(x_in) if args.gpu >= 0: x_in.to_gpu() st = time.time() for i in range(10): z = enc(x_in) x_out = dec(z) ts = time.time() - st print("time:{:.2f}".format(ts/10.0)) if args.gpu >= 0: out_ary = x_out.data.get()[0] else: out_ary = x_out.data[0] #img_show = np.zeros((inimg.shape[0], inimg.shape[1], inimg.shape[2]*2)) #img_show[:,:,:inimg.shape[2]] = inimg #img_show[:,:outimg.shape[1],inimg.shape[2]:inimg.shape[2]+outimg.shape[2]] = outimg outimg = out_ary[:,0:h,0:w] # trim paddings img_show = np.concatenate((inimg, outimg), axis=2) bgrpic = to_bgr(img_show).copy() cv2.putText(bgrpic,"input",(3,15),cv2.FONT_HERSHEY_DUPLEX, 0.5,(255,0,0)) cv2.putText(bgrpic,"output",(w+3,15),cv2.FONT_HERSHEY_DUPLEX, 0.5,(255,0,0)) cv2.imshow("result", bgrpic) cv2.waitKey()
def main(): parser = argparse.ArgumentParser(description='chainer implementation of pix2pix') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--seed', type=int, default=0, help='Random seed') parser.add_argument('--model', '-m', default='', help='model snapshot') parser.add_argument('--enc', '-e', type=str, default='enc_iter_60000.npz', help='encoder snapshot') parser.add_argument('--dec', '-d', type=str, default='dec_iter_60000.npz', help='decoder snapshot') parser.add_argument('--out', '-o', type=str, default='out', help='output dir') parser.add_argument('--input', '-i', default='sample.jpg', help='input jpg', required=True) parser.add_argument('--contour', '-c', action='store_true', help='from contour image or not') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) enc = Encoder(in_ch=3) dec = Decoder(out_ch=3) dis = Discriminator(in_ch=3, out_ch=3) if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current enc.to_gpu() # Copy the model to the GPU dec.to_gpu() dis.to_gpu() # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec') return optimizer if args.model: opt_enc = make_optimizer(enc) opt_dec = make_optimizer(dec) opt_dis = make_optimizer(dis) # Set up a trainer updater = FacadeUpdater( models=(enc, dec, dis), iterator={}, optimizer={ 'enc': opt_enc, 'dec': opt_dec, 'dis': opt_dis}, device=args.gpu) trainer = training.Trainer(updater, (200, 'epoch'), out='generate/') chainer.serializers.load_npz(args.model, trainer) elif args.enc and args.dec: chainer.serializers.load_npz(args.enc, enc) chainer.serializers.load_npz(args.dec, dec) if not args.contour: from make_contour import get_contour_image get_contour_image(args.input) generate_image_from_contour(args.input, enc, dec, args.out)
def main(): parser = argparse.ArgumentParser( description='chainer implementation of pix2pix') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--dec', '-d', type=str, default='dec_model.npz', help='decoder model') parser.add_argument('--enc', '-e', type=str, default='enc_model.npz', help='encoder model') args = parser.parse_args() enc = Encoder(in_ch=3) dec = Decoder(out_ch=1) chainer.serializers.load_npz(str(MODEL_PATH.joinpath(args.dec)), dec) chainer.serializers.load_npz(str(MODEL_PATH.joinpath(args.enc)), enc) if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current enc.to_gpu() # Copy the model to the GPU dec.to_gpu() xp = enc.xp branch_images = np.zeros((12, 256, 256), dtype=np.uint8) plant_images = np.zeros((12, 256, 256, 3), dtype=np.uint8) p2b_images = np.zeros((12, 256, 256), dtype=np.uint8) for i in tqdm(range(1,12)): branch_path = DATASET_PATH.joinpath('branch', str(i)) plant_path = DATASET_PATH.joinpath('plant', str(i)) name = random.choice([_ for _ in branch_path.glob('*.png')]).name branch_image_path = branch_path.joinpath(name) plant_image_path = plant_path.joinpath(name) # open image branch_image = np.asarray(Image.open(branch_image_path).convert('L')) branch_images[i-1,:] = branch_image plant_image = np.asarray(Image.open(plant_image_path).convert('RGB')) plant_images[i-1,:] = plant_image plant_image = xp.array(plant_image).astype("f").transpose(2, 0, 1) / 128.0-1.0 plant_image = plant_image.reshape(1, *plant_image.shape) with chainer.no_backprop_mode(), chainer.using_config('train', False): p2b_image = np.asarray(dec(enc(plant_image)).data.get()) p2b_image = np.asarray(np.clip(p2b_image * 128 + 128, 0.0, 255.0), dtype=np.uint8).reshape(256, 256) p2b_images[i-1, :] = p2b_image Image.fromarray(branch_images.reshape(3, 4, 256, 256).transpose(0, 2, 1, 3).reshape(3*256, 4*256))\ .save(str(RESULT_PATH.joinpath('branch_image.png'))) Image.fromarray(plant_images.reshape(3, 4, 256, 256, 3).transpose(0, 2, 1, 3, 4).reshape(3*256, 4*256, 3))\ .save(str(RESULT_PATH.joinpath('plant_image.png'))) Image.fromarray(p2b_images.reshape(3, 4, 256, 256).transpose(0, 2, 1, 3).reshape(3*256, 4*256))\ .save(str(RESULT_PATH.joinpath('p2b_image.png')))
def experiment(): E_net = Encoder().double() T_net = Transition().double() optimizer_predict = torch.optim.RMSprop(itertools.chain( E_net.parameters(), T_net.parameters()), lr=HP_DICT['learning_rate']) data = RectsData(HP_DICT) if HP_DICT['GPU']: E_net = E_net.to('cuda') T_net = T_net.to('cuda') data = data.to('cuda') plt.ion() return train(E_net, T_net, data, optimizer_predict)
def encode(args): enc = Encoder() npz.load_npz(args.enc, enc) if args.gpu >= 0: chainer.cuda.get_device_from_id(args.gpu).use() enc.to_gpu() xp = enc.xp image = CelebADataset([args.infile])[0] x = Variable(xp.asarray([image])) / 255. x = F.resize_images(x, (64, 64)) with chainer.using_config('train', False): z = enc(x) return z, x.data[0]
def main(): parser = argparse.ArgumentParser( description='chainer implementation of pix2pix') parser.add_argument('--batchsize', '-b', type=int, default=1, help='Number of images in each mini-batch') parser.add_argument('--epoch', '-e', type=int, default=200, help='Number of sweeps over the dataset to train') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--out', '-o', default='result', help='Directory to output the result') parser.add_argument('--enc-npz', default='', required=True) parser.add_argument('--dec-npz', default='', required=True) parser.add_argument('input') args = parser.parse_args() # Set up a neural network to train enc = Encoder(in_ch=3) dec = Decoder(out_ch=3) chainer.serializers.load_npz(args.enc_npz, enc) chainer.serializers.load_npz(args.dec_npz, dec) # load image and predict img = Image.open(args.input).convert('RGB') img = np.asarray(img).astype(np.float32) img /= 255.0 img = img.transpose(2, 0, 1).reshape((1, 3, 256, 256)) x = chainer.Variable(img) z = enc(x) y = dec(z) # (batchsize, ch, width, height) y = y.data.reshape(3, 256, 256).transpose(1, 2, 0).clip(0, 1) import matplotlib.pyplot as plt plt.imshow(y) plt.show() if False: import scipy.misc scipy.misc.imsave("output.png", y)
def main(): parser = argparse.ArgumentParser( description='chainer implementation of pix2pix') parser.add_argument('--batchsize', '-b', type=int, default=1, help='Number of images in each mini-batch') parser.add_argument('--epoch', '-e', type=int, default=200, help='Number of sweeps over the dataset to train') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--out', '-o', default='result', help='Directory to output the result') parser.add_argument('--enc-npz', default='', required=True) parser.add_argument('--dec-npz', default='', required=True) args = parser.parse_args() # Set up a neural network to train enc = Encoder(in_ch=3) dec = Decoder(out_ch=3) chainer.serializers.load_npz(args.enc_npz, enc) chainer.serializers.load_npz(args.dec_npz, dec) # graph input = np.zeros((1, 3, 256, 256), dtype=np.float32) x = chainer.Variable(input) z = enc(x) y = dec(z) graph = ChainerConverter().convert([x], [y]) exec_info = generate_descriptor("webassembly", graph) exec_info.save(args.out)
def test_random(): enc = Encoder((None, ) + input_size, vgg_param_file) dec = Decoder(enc.output_shape, dec_param_file) X = T.tensor4('input') Y = T.tensor4('style') weights = [T.vector('weights') for i in range(len(vgg_info))] nn.set_training_off() X_styled = dec(enc((X, Y), weights)) test = nn.function([X, Y] + weights, X_styled, name='test generator') style_folder = os.listdir(style_img_folder) input_folder = os.listdir(input_img_folder) time_list = [] if not os.path.exists('outputs'): os.mkdir('outputs') for style_file in style_folder: sty_img = prep_image_test( misc.imread(style_img_folder + '/' + style_file)) for input_file in input_folder: try: input_img = prep_image_test( misc.imread(input_img_folder + '/' + input_file)) except ValueError: continue for i in range(num_styles): start = time.time() output = test(input_img, sty_img, *get_weights(vgg_info)) time_list.append(time.time() - start) output = np.transpose(output[0], (1, 2, 0)) misc.imsave( os.path.join( 'outputs', input_file[:-4] + '_' + style_file[:-4] + '_%d.jpg' % i), output) print('Took %f s/image' % np.mean(time_list)) print('Testing finished!')
def __init__(self): super(Agent, self).__init__() #self.config = config self.action = Action() self.dataset = Dataset() self.device = torch.device( "cuda" if torch.cuda.is_available() else "cpu") self.grad_clip = 5 embed_dim = 512 # dimension of word embeddings decoder_dim = 512 # dimension of decoder RNN dropout = 0.5 #batch_size = 32 self.n_epochs = 20 self.encoder = Encoder() self.decoder = Decoder(embed_dim=embed_dim, decoder_dim=decoder_dim, vocab_size=self.dataset.get_word_map_len(), dropout=dropout)
def main(): parser = argparse.ArgumentParser( description='chainer creating pictures of seat') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--data', '-d', default='./detect', help='Directory of image files.') parser.add_argument('--mask', '-ma', default='./mask', help='Directory of image files.') parser.add_argument('--out', '-o', default='./result', help='Directory to output the result') parser.add_argument('--seed', type=int, default=0, help='Random seed') parser.add_argument('--model', '-m', default='snapshot_iter_83000.npz', help='Loading model') parser.add_argument('--batchsize', '-b', default=16, type=int, help='The same value as that of trainer') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) print('# Minibatch-size: {}'.format(args.batchsize)) print('') # Set up a neural network to train enc = Encoder(in_ch=4) dec = Decoder(out_ch=3) if args.gpu >= 0: cuda.get_device(args.gpu).use() # Make a specified GPU current enc.to_gpu() # Copy the model to the GPU dec.to_gpu() xp = cuda.cupy # Setup an model model = args.model print('---Loading model---') serializers.load_npz(model, enc, path='updater/model:enc/') serializers.load_npz(model, dec, path='updater/model:dec/') # Trun off a trainer chainer.config.train = False chainer.config.debug = False chainer.config.enable_backprop = False print('Setuped model!') # Load datasets print('---Loading datasets---') data_path = args.data data_sets = glob.glob(data_path + '/*.jpg') data_mask = glob.glob(args.mask + '/*.jpg') dataset = [] names = [] for data in data_sets: d_name = os.path.basename(data) d_name = d_name[:-4] img = Image.open(data) img = xp.asarray(img).transpose(2, 0, 1) for _ in range(10): mask = random.choice(data_mask) mask = Image.open(mask) mask = xp.asarray(mask) mask = mask[xp.newaxis, :, :] img_ = img + mask img_ = xp.asarray(img_).astype('f') / 128.0 - 1.0 mask = xp.asarray(mask).astype('f') / 128.0 - 1.0 img_ = xp.concatenate([img_, mask], axis=0) dataset.append(img_) f_name = d_name + '_' + str(_) names.append(f_name) print('Setuped datasets!') # Create picture print('---Creating---') in_ch = 4 in_h = 256 in_w = 256 out_put = 0 batch_size = args.batchsize out_dir = args.out if not (os.path.exists(out_dir)): os.mkdir(out_dir) _ = 0 for name, data in zip(names, dataset): X_in = xp.zeros((batch_size, in_ch, in_h, in_w)).astype("f") for i in range(batch_size): X_in[i, :] = xp.asarray(data) X_in = Variable(X_in) z = enc(X_in) X_out = dec(z) out_put = xp.asarray(X_out.data) out_put = out_put[0] out_put += 1.0 out_put *= 128.0 xp.save(out_dir + '/' + name, out_put) _ += 1 print('created {} / {}'.format(_, len(dataset))) print('Finished all process!') print('Numpy shape : ', out_put.shape) print('Number of Numpy file : ', len(dataset))
def main(): parser = argparse.ArgumentParser(description='Train Encoder') parser.add_argument('--batchsize', '-b', type=int, default=64, help='Number of images in each mini-batch') parser.add_argument('--epoch', '-e', type=int, default=100, help='Number of sweeps over the dataset to train') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--dataset', '-i', default='data/celebA/', help='Directory of image files.') parser.add_argument('--out', '-o', default='result', help='Directory to output the result') parser.add_argument('--resume', '-r', default='', help='Resume the training from snapshot') parser.add_argument('--snapshot_interval', type=int, default=10000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=1000, help='Interval of displaying log to console') parser.add_argument('--gen', default='gen.npz') parser.add_argument('--enc', default=None) args = parser.parse_args() print('GPU: {}'.format(args.gpu)) print('# batchsize: {}'.format(args.batchsize)) print('# epoch: {}'.format(args.epoch)) print('') # Set up a neural network to train gen = Generator() npz.load_npz(args.gen, gen) enc = Encoder() if args.enc is not None: npz.load_npz(args.enc, enc) if args.gpu >= 0: chainer.cuda.get_device_from_id(args.gpu).use() gen.to_gpu() enc.to_gpu() # Setup an optimizer def make_optimizer(model, alpha=0.0005, beta1=0.9): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.0001), 'hook_dec') return optimizer opt_gen = make_optimizer(gen) gen.disable_update() opt_enc = make_optimizer(enc) # Setup a dataset all_files = os.listdir(args.dataset) image_files = [f for f in all_files if ('png' in f or 'jpg' in f)] print('{} contains {} image files'.format(args.dataset, len(image_files))) train = CelebADataset(paths=image_files, root=args.dataset) train_iter = chainer.iterators.SerialIterator(train, args.batchsize) # Set up a trainer updater = EncUpdater( models=(gen, enc), iterator=train_iter, optimizer={'gen': opt_gen, 'enc': opt_enc}, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) snapshot_interval = (args.snapshot_interval, 'iteration') display_interval = (args.display_interval, 'iteration') trainer.extend( extensions.snapshot(filename='snapshot_enc_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.ExponentialShift( 'alpha', 0.5, optimizer=opt_enc), trigger=(10, 'epoch')) trainer.extend(extensions.snapshot_object( enc, 'enc_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.LogReport(trigger=display_interval, log_name='train_enc.log')) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'enc/loss', ]), trigger=display_interval) trainer.extend(extensions.PlotReport( ['enc/loss'], trigger=display_interval, file_name='enc-loss.png')) trainer.extend(extensions.ProgressBar(update_interval=10)) if args.resume: chainer.serializers.load_npz(args.resume, trainer) # Run the training trainer.run()
help='Output file name') args = parser.parse_args() if args.use_cuda and not torch.cuda.is_available(): raise Exception('No GPU found, please run without --cuda') device = torch.device('cuda' if args.use_cuda else 'cpu') dataset = DatasetFromFolder(args.input, 'file', params.sr, params.length, params.frame_length, params.hop, params.n_mels, 'valid', None) data_loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=1) encoder = Encoder(params.upscale_factors, params.n_wavenets * params.n_layers * params.n_loops, params.r, params.n_mels) wavenet = UniWaveNet(params.n_wavenets, params.n_layers, params.n_loops, params.a, params.r, params.s) encoder = encoder.to(device) wavenet = wavenet.to(device) encoder.load_state_dict(torch.load(args.encoder_model)) wavenet.load_state_dict(torch.load(args.wavenet_model)) print(torch.load(args.wavenet_model)['wavenet_list.0.conv_in.weight'][0][0][0]) print(torch.load(args.wavenet_model)['wavenet_list.0.conv1x1.weight'][0][0][0]) print( torch.load(args.wavenet_model)['wavenet_list.0.conv_out.weight'][0][0][0]) print(torch.load(args.wavenet_model)['wavenet_list.1.conv_in.weight'][0][0][0])
x = -np.log(1 / ((x + 1) / 2) - 1) return x def _istft(y): _, hop_length, win_length = _stft_parameters() return librosa.istft(y, hop_length=hop_length, win_length=win_length) def save_wav(wav, path): wav *= 1 / max(0.01, np.max(np.abs(wav))) librosa.output.write_wav(path, wav, hparams.sample_rate) # Set up a neural network enc = Encoder(in_ch=2) dec = Decoder(out_ch=2) # load terained models chainer.serializers.load_npz(MODEL_DIR + "enc_iter_%d.npz" % ITER, enc) chainer.serializers.load_npz(MODEL_DIR + "dec_iter_%d.npz" % ITER, dec) #img = Image.open('../input/png/%d_ritsuko.png'%MUSIC_IDX) #img = np.asarray(img).astype("f")/128.0-1.0 img = sp2.astype("f") / 128.0 - 1.0 h, w, c = img.shape print(img.shape) img = img.transpose(2, 0, 1) img = img[:2, :, :] img = img.reshape(1, 2, h, w)
def main(): attr_columns = [ '5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings', 'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young' ] parser = argparse.ArgumentParser(description='Get Attribute Vector') parser.add_argument('--batchsize', '-b', type=int, default=512, help='Number of images in each mini-batch') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--dataset', '-i', default='data/celebA/', help='Directory of image files.') parser.add_argument('--attr_list', '-a', default='data/list_attr_celeba.txt') parser.add_argument('--get_attr', default='all', nargs='+', choices=attr_columns + ['all']) parser.add_argument('--outfile', '-o', default='attr_vec.json') parser.add_argument('--enc', default='pre-trained/enc_iter_310000.npz') args = parser.parse_args() enc = Encoder() npz.load_npz(args.enc, enc) if args.gpu >= 0: chainer.cuda.get_device_from_id(args.gpu).use() enc.to_gpu() all_files = os.listdir(args.dataset) image_files = [f for f in all_files if ('png' in f or 'jpg' in f)] vectors = {} attr_df = pd.read_csv(args.attr_list, delim_whitespace=True, header=1) if args.get_attr == 'all': args.get_attr = attr_columns for attr_name in tqdm(list(set(args.get_attr) & set(attr_df.columns))): with_attr_files = attr_df[attr_df[attr_name] == 1].index.tolist() with_attr_files = list(set(with_attr_files) & set(image_files)) with_attr_vec = get_vector(enc, with_attr_files, args) without_attr_files = attr_df[attr_df[attr_name] != 1].index.tolist() without_attr_files = list(set(without_attr_files) & set(image_files)) without_attr_vec = get_vector(enc, without_attr_files, args) vectors[attr_name] = (with_attr_vec - without_attr_vec).tolist() with open(args.outfile, 'w') as f: f.write( json.dumps(vectors, indent=4, sort_keys=True, separators=(',', ': ')))
def train(folding_id, inliner_classes, ic): cfg = get_cfg_defaults() cfg.merge_from_file('configs/mnist.yaml') cfg.freeze() logger = logging.getLogger("logger") zsize = cfg.MODEL.LATENT_SIZE output_folder = os.path.join('results_' + str(folding_id) + "_" + "_".join([str(x) for x in inliner_classes])) os.makedirs(output_folder, exist_ok=True) os.makedirs('models', exist_ok=True) train_set, _, _ = make_datasets(cfg, folding_id, inliner_classes) logger.info("Train set size: %d" % len(train_set)) G = Generator(cfg.MODEL.LATENT_SIZE, channels=cfg.MODEL.INPUT_IMAGE_CHANNELS) G.weight_init(mean=0, std=0.02) D = Discriminator(channels=cfg.MODEL.INPUT_IMAGE_CHANNELS) D.weight_init(mean=0, std=0.02) E = Encoder(cfg.MODEL.LATENT_SIZE, channels=cfg.MODEL.INPUT_IMAGE_CHANNELS) E.weight_init(mean=0, std=0.02) if cfg.MODEL.Z_DISCRIMINATOR_CROSS_BATCH: ZD = ZDiscriminator_mergebatch(zsize, cfg.TRAIN.BATCH_SIZE) else: ZD = ZDiscriminator(zsize, cfg.TRAIN.BATCH_SIZE) ZD.weight_init(mean=0, std=0.02) lr = cfg.TRAIN.BASE_LEARNING_RATE G_optimizer = optim.Adam(G.parameters(), lr=lr, betas=(0.5, 0.999)) D_optimizer = optim.Adam(D.parameters(), lr=lr, betas=(0.5, 0.999)) GE_optimizer = optim.Adam(list(E.parameters()) + list(G.parameters()), lr=lr, betas=(0.5, 0.999)) ZD_optimizer = optim.Adam(ZD.parameters(), lr=lr, betas=(0.5, 0.999)) BCE_loss = nn.BCELoss() sample = torch.randn(64, zsize).view(-1, zsize, 1, 1) tracker = LossTracker(output_folder=output_folder) for epoch in range(cfg.TRAIN.EPOCH_COUNT): G.train() D.train() E.train() ZD.train() epoch_start_time = time.time() data_loader = make_dataloader(train_set, cfg.TRAIN.BATCH_SIZE, torch.cuda.current_device()) train_set.shuffle() if (epoch + 1) % 30 == 0: G_optimizer.param_groups[0]['lr'] /= 4 D_optimizer.param_groups[0]['lr'] /= 4 GE_optimizer.param_groups[0]['lr'] /= 4 ZD_optimizer.param_groups[0]['lr'] /= 4 print("learning rate change!") for y, x in data_loader: x = x.view(-1, cfg.MODEL.INPUT_IMAGE_CHANNELS, cfg.MODEL.INPUT_IMAGE_SIZE, cfg.MODEL.INPUT_IMAGE_SIZE) y_real_ = torch.ones(x.shape[0]) y_fake_ = torch.zeros(x.shape[0]) y_real_z = torch.ones( 1 if cfg.MODEL.Z_DISCRIMINATOR_CROSS_BATCH else x.shape[0]) y_fake_z = torch.zeros( 1 if cfg.MODEL.Z_DISCRIMINATOR_CROSS_BATCH else x.shape[0]) ############################################# D.zero_grad() D_result = D(x).squeeze() D_real_loss = BCE_loss(D_result, y_real_) z = torch.randn((x.shape[0], zsize)).view(-1, zsize, 1, 1) z = Variable(z) x_fake = G(z).detach() D_result = D(x_fake).squeeze() D_fake_loss = BCE_loss(D_result, y_fake_) D_train_loss = D_real_loss + D_fake_loss D_train_loss.backward() D_optimizer.step() tracker.update(dict(D=D_train_loss)) ############################################# G.zero_grad() z = torch.randn((x.shape[0], zsize)).view(-1, zsize, 1, 1) z = Variable(z) x_fake = G(z) D_result = D(x_fake).squeeze() G_train_loss = BCE_loss(D_result, y_real_) G_train_loss.backward() G_optimizer.step() tracker.update(dict(G=G_train_loss)) ############################################# ZD.zero_grad() z = torch.randn((x.shape[0], zsize)).view(-1, zsize) z = Variable(z) ZD_result = ZD(z).squeeze() ZD_real_loss = BCE_loss(ZD_result, y_real_z) z = E(x).squeeze().detach() ZD_result = ZD(z).squeeze() ZD_fake_loss = BCE_loss(ZD_result, y_fake_z) ZD_train_loss = ZD_real_loss + ZD_fake_loss ZD_train_loss.backward() ZD_optimizer.step() tracker.update(dict(ZD=ZD_train_loss)) # ############################################# E.zero_grad() G.zero_grad() z = E(x) x_d = G(z) ZD_result = ZD(z.squeeze()).squeeze() E_train_loss = BCE_loss(ZD_result, y_real_z) * 1.0 Recon_loss = F.binary_cross_entropy(x_d, x.detach()) * 2.0 (Recon_loss + E_train_loss).backward() GE_optimizer.step() tracker.update(dict(GE=Recon_loss, E=E_train_loss)) # ############################################# comparison = torch.cat([x, x_d]) save_image(comparison.cpu(), os.path.join(output_folder, 'reconstruction_' + str(epoch) + '.png'), nrow=x.shape[0]) epoch_end_time = time.time() per_epoch_ptime = epoch_end_time - epoch_start_time logger.info( '[%d/%d] - ptime: %.2f, %s' % ((epoch + 1), cfg.TRAIN.EPOCH_COUNT, per_epoch_ptime, tracker)) tracker.register_means(epoch) tracker.plot() with torch.no_grad(): resultsample = G(sample).cpu() save_image( resultsample.view(64, cfg.MODEL.INPUT_IMAGE_CHANNELS, cfg.MODEL.INPUT_IMAGE_SIZE, cfg.MODEL.INPUT_IMAGE_SIZE), os.path.join(output_folder, 'sample_' + str(epoch) + '.png')) logger.info("Training finish!... save training results") os.makedirs("models", exist_ok=True) print("Training finish!... save training results") torch.save(G.state_dict(), "models/Gmodel_%d_%d.pkl" % (folding_id, ic)) torch.save(E.state_dict(), "models/Emodel_%d_%d.pkl" % (folding_id, ic))
def main(): parser = argparse.ArgumentParser( description="chainer implementation of Unet") parser.add_argument("--batchsize", "-b", type=int, default=1, help="Number of images in each mini-batch") parser.add_argument("--epoch", "-e", type=int, default=200, help="epoch") parser.add_argument("--gpu", "-g", type=int, default=-1, help="GPU ID") parser.add_argument("--dataset", "-i", default="./train/", help="Directory of image files") parser.add_argument("--out", "-o", default="result/", help="Directory to output the result") parser.add_argument('--resume', '-r', default='', help='Resume the training from snapshot') parser.add_argument('--seed', type=int, default=0, help='Random seed') parser.add_argument('--snapshot_interval', type=int, default=1000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=100, help='Interval of displaying log to console') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) print('# Minibatch-size: {}'.format(args.batchsize)) print('# epoch: {}'.format(args.epoch)) print('') # Only Encoder Decoder with Unet enc = Encoder(in_ch=3) # in_ch => 3(YCbCr) dec = Decoder(out_ch=3) # out_ch => 3(DCT) # GPU set up if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() enc.to_gpu() dec.to_gpu() # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec') return optimizer opt_enc = make_optimizer(enc) opt_dec = make_optimizer(dec) train_d = ImgDCTDataset(args.dataset, data_range=(0, 1000)) test_d = ImgDCTDataset(args.dataset, data_range=(1000, 2000)) train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize) test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize) updater = FacadeUpdater(models=(enc, dec), iterator={ 'main': train_iter, 'test': test_iter }, optimizer={ 'enc': opt_enc, 'dec': opt_dec }, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) snapshot_interval = (args.snapshot_interval, 'iteration') display_interval = (args.display_interval, 'iteration') trainer.extend( extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( enc, 'enc_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dec, 'dec_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) # trainer.extend(extensions.snapshot_object( # dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'enc/loss', 'dec/loss', ]), trigger=display_interval) trainer.extend(extensions.ProgressBar(update_interval=10)) if args.resume: # Resume from a snapshot chainer.serializers.load_npz(args.resume, trainer) # Run the training trainer.run()
def main(): parser = argparse.ArgumentParser(description='chainer implementation of pix2pix') parser.add_argument('--batchsize', '-b', type=int, default=1, help='Number of images in each mini-batch') parser.add_argument('--epoch', '-e', type=int, default=200, help='Number of sweeps over the dataset to train') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--dataset', '-i', default='./facade/base', help='Directory of image files.') parser.add_argument('--out', '-o', default='result', help='Directory to output the result') parser.add_argument('--resume', '-r', default='', help='Resume the training from snapshot') parser.add_argument('--seed', type=int, default=0, help='Random seed') parser.add_argument('--snapshot_interval', type=int, default=1000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=100, help='Interval of displaying log to console') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) print('# Minibatch-size: {}'.format(args.batchsize)) print('# epoch: {}'.format(args.epoch)) print('') # Set up a neural network to train enc = Encoder(in_ch=12) dec = Decoder(out_ch=3) dis = Discriminator(in_ch=12, out_ch=3) if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current enc.to_gpu() # Copy the model to the GPU dec.to_gpu() dis.to_gpu() # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec') return optimizer opt_enc = make_optimizer(enc) opt_dec = make_optimizer(dec) opt_dis = make_optimizer(dis) train_d = FacadeDataset(args.dataset, data_range=(1,300)) test_d = FacadeDataset(args.dataset, data_range=(300,379)) #train_iter = chainer.iterators.MultiprocessIterator(train_d, args.batchsize, n_processes=4) #test_iter = chainer.iterators.MultiprocessIterator(test_d, args.batchsize, n_processes=4) train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize) test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize) # Set up a trainer updater = FacadeUpdater( models=(enc, dec, dis), iterator={ 'main': train_iter, 'test': test_iter}, optimizer={ 'enc': opt_enc, 'dec': opt_dec, 'dis': opt_dis}, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) snapshot_interval = (args.snapshot_interval, 'iteration') display_interval = (args.display_interval, 'iteration') trainer.extend(extensions.snapshot( filename='snapshot_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( enc, 'enc_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dec, 'dec_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'enc/loss', 'dec/loss', 'dis/loss', ]), trigger=display_interval) trainer.extend(extensions.ProgressBar(update_interval=10)) trainer.extend( out_image( updater, enc, dec, 5, 5, args.seed, args.out), trigger=snapshot_interval) if args.resume: # Resume from a snapshot chainer.serializers.load_npz(args.resume, trainer) # Run the training trainer.run()
help='Resume the training from snapshot') parser.add_argument('--seed', type=int, default=0, help='Random seed') parser.add_argument('--snapshot_interval', type=int, default=1000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=100, help='Interval of displaying log to console') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) print('# Minibatch-size: {}'.format(args.batchsize)) print('# epoch: {}'.format(args.epoch)) print('') # Set up a neural network to train enc = Encoder(in_ch=5) dec = Decoder(out_ch=3) dis = Discriminator(in_ch=5, out_ch=3) if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current enc.to_gpu() # Copy the model to the GPU dec.to_gpu() dis.to_gpu() # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec') return optimizer
def main(): parser = argparse.ArgumentParser( description="chainer implementation of pix2pix") parser.add_argument("--batchsize", "-b", type=int, default=1, help="Number of images in each mini-batch") parser.add_argument("--epoch", "-e", type=int, default=40000, help="Number of sweeps over the dataset to train") parser.add_argument("--gpu", "-g", type=int, default=-1, help="GPU ID (negative value indicates CPU)") parser.add_argument("--dataset", "-i", default="./input/png/", help="Directory of image files.") parser.add_argument("--out", "-o", default="D:/output/imasUtaConverter/", help="Directory to output the result") parser.add_argument("--resume", "-r", default="", help="Resume the training from snapshot") parser.add_argument("--seed", type=int, default=0, help="Random seed") parser.add_argument("--snapshot_interval", type=int, default=10000, help="Interval of snapshot") parser.add_argument("--display_interval", type=int, default=20, help="Interval of displaying log to console") args = parser.parse_args() print("GPU: {}".format(args.gpu)) print("# Minibatch-size: {}".format(args.batchsize)) print("# epoch: {}".format(args.epoch)) print("") # Set up a neural network to train enc = Encoder(in_ch=2) dec = Decoder(out_ch=2) dis = Discriminator(in_ch=2, out_ch=2) if args.gpu >= 0: chainer.backends.cuda.get_device_from_id( args.gpu).use() # Make a specified GPU current enc.to_gpu() # Copy the model to the GPU dec.to_gpu() dis.to_gpu() # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), "hook_dec") return optimizer opt_enc = make_optimizer(enc) opt_dec = make_optimizer(dec) opt_dis = make_optimizer(dis) train_d = FacadeDataset(args.dataset, data_range=(0, 38)) test_d = FacadeDataset(args.dataset, data_range=(38, 40)) #train_iter = chainer.iterators.MultiprocessIterator(train_d, args.batchsize, n_processes=4) #test_iter = chainer.iterators.MultiprocessIterator(test_d, args.batchsize, n_processes=4) train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize) test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize) # Set up a trainer updater = FacadeUpdater(models=(enc, dec, dis), iterator={ "main": train_iter, "test": test_iter }, optimizer={ "enc": opt_enc, "dec": opt_dec, "dis": opt_dis }, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, "epoch"), out=args.out) snapshot_interval = (args.snapshot_interval, "iteration") display_interval = (args.display_interval, "iteration") trainer.extend( extensions.snapshot(filename="snapshot_iter_{.updater.iteration}.npz"), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( enc, "enc_iter_{.updater.iteration}.npz"), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dec, "dec_iter_{.updater.iteration}.npz"), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dis, "dis_iter_{.updater.iteration}.npz"), trigger=snapshot_interval) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ "epoch", "iteration", "enc/loss", "dec/loss", "dis/loss", ]), trigger=display_interval) #trainer.extend(extensions.PlotReport(["enc/loss", "dis/loss"], x_key="epoch", file_name="loss.png")) trainer.extend(extensions.ProgressBar(update_interval=20)) trainer.extend(out_image(updater, enc, dec, 1, 10, args.seed, args.out), trigger=snapshot_interval) if args.resume: # Resume from a snapshot chainer.serializers.load_npz(args.resume, trainer) # Run the training trainer.run()
def main(): parser = argparse.ArgumentParser( description='chainer implementation of pix2pix') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--seed', type=int, default=0, help='Random seed') parser.add_argument('--model', '-m', default='', help='model snapshot') parser.add_argument('--input', '-i', default='../images/generate/sample.jpg', help='input jpg') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) # Set up a neural network to train enc = Encoder(in_ch=3) dec = Decoder(out_ch=3) dis = Discriminator(in_ch=3, out_ch=3) if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current enc.to_gpu() # Copy the model to the GPU dec.to_gpu() dis.to_gpu() # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec') return optimizer opt_enc = make_optimizer(enc) opt_dec = make_optimizer(dec) opt_dis = make_optimizer(dis) if os.path.exists('generate_tmp'): shutil.rmtree('generate_tmp') os.mkdir('generate_tmp') os.mkdir('generate_tmp/base') os.mkdir('generate_tmp/label') shutil.copyfile(args.input, 'generate_tmp/base/tmp.jpg') shutil.copyfile(args.input, 'generate_tmp/label/tmp.jpg') test_d = FacadeDataset('generate_tmp/') test_iter = chainer.iterators.SerialIterator(test_d, 1) # Set up a trainer updater = FacadeUpdater(models=(enc, dec, dis), iterator={}, optimizer={ 'enc': opt_enc, 'dec': opt_dec, 'dis': opt_dis }, device=args.gpu) trainer = training.Trainer(updater, (200, 'epoch'), out='../results/generate/') chainer.serializers.load_npz(args.model, trainer) out_image(updater, enc, dec, 1, 1, args.seed, '../results/generate/', True, test_iter)(trainer)
params.seed) valid_dataset = DatasetFromFolder( params.root, params.dataset_type, params.sr, params.length, params.frame_length, params.hop, params.n_mels, 'valid', params.seed) train_data_loader = torch.utils.data.DataLoader( dataset=train_dataset, batch_size=params.batch_size, shuffle=True) valid_data_loader = torch.utils.data.DataLoader( dataset=valid_dataset, batch_size=params.batch_size, shuffle=False) encoder = Encoder( params.upscale_factors, params.n_wavenets * params.n_layers * params.n_loops, params.r, params.n_mels).to(device) wavenet = UniWaveNet( params.n_wavenets, params.n_layers, params.n_loops, params.a, params.r, params.s).to(device) optimizer = torch.optim.Adam( list(wavenet.parameters()) + list(encoder.parameters()), lr=params.lr) train_writer = tensorboardX.SummaryWriter( str(pathlib.Path(params.output_dir, 'train'))) valid_writer = tensorboardX.SummaryWriter( str(pathlib.Path(params.output_dir, 'valid'))) trainer = UniWaveNetTrainer( train_data_loader, valid_data_loader, train_writer, valid_writer,
def main(): parser = argparse.ArgumentParser( description='chainer implementation of pix2pix') parser.add_argument('--batchsize', '-b', type=int, default=1, help='Number of images in each mini-batch') parser.add_argument('--epoch', '-e', type=int, default=200, help='Number of sweeps over the dataset to train') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--dataset', '-i', default='./facade/base', help='Directory of image files.') parser.add_argument('--out', '-o', default='result', help='Directory to output the result') parser.add_argument('--resume', '-r', default='', help='Resume the training from snapshot') parser.add_argument('--seed', type=int, default=0, help='Random seed') parser.add_argument('--snapshot_interval', type=int, default=1000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=100, help='Interval of displaying log to console') parser.add_argument( '--n_processes', type=int, default=None, help='processes of chainer.iterators.MultiprocessIterator') parser.add_argument( '--shared_mem', type=int, default=None, help= 'shared memory per data, for chainer.iterators.MultiprocessIterator. None means auto ajust.' ) parser.add_argument('--audio_dataset_second', type=int, default=None, help='time length(second) of train audio data .') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) print('# Minibatch-size: {}'.format(args.batchsize)) print('# epoch: {}'.format(args.epoch)) print('') args.out = os.path.join(args.out, datetime.now().strftime("%Y%m%d_%H%M%S")) util.audio_dataset_second = args.audio_dataset_second if args.batchsize > 1: assert util.audio_dataset_second != None, "when minibatch training (e.g. --batchsize > 1), --audio_dataset_second option is required." # Set up a neural network to train enc = Encoder(in_ch=2) dec = Decoder(out_ch=2) dis = Discriminator(in_ch=2, out_ch=2) if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current enc.to_gpu() # Copy the model to the GPU dec.to_gpu() dis.to_gpu() # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec') return optimizer opt_enc = make_optimizer(enc) opt_dec = make_optimizer(dec) opt_dis = make_optimizer(dis) train_d = Vp2pDataset(args.dataset + "/train") test_d = Vp2pDataset(args.dataset + "/test") train_iter = chainer.iterators.MultiprocessIterator( train_d, args.batchsize, n_processes=args.n_processes, shared_mem=args.shared_mem) test_iter = chainer.iterators.MultiprocessIterator( test_d, args.batchsize, n_processes=args.n_processes, shared_mem=args.shared_mem) # train_iter = chainer.iterators.MultiprocessIterator(train_d, args.batchsize) # test_iter = chainer.iterators.MultiprocessIterator(test_d, args.batchsize) # Set up a trainer updater = VoiceP2PUpdater(models=(enc, dec, dis), iterator={ 'main': train_iter, 'test': test_iter }, optimizer={ 'enc': opt_enc, 'dec': opt_dec, 'dis': opt_dis }, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) snapshot_interval = (args.snapshot_interval, 'iteration') display_interval = (args.display_interval, 'iteration') trainer.extend( extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( enc, 'enc_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dec, 'dec_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'enc/loss', 'dec/loss', 'dis/loss', ]), trigger=display_interval) trainer.extend(extensions.ProgressBar(update_interval=10)) trainer.extend(out_image(updater, enc, dec, 5, 5, args.seed, args.out), trigger=snapshot_interval) if args.resume: # Resume from a snapshot chainer.serializers.load_npz(args.resume, trainer) # Run the training trainer.run()
def main(): parser = argparse.ArgumentParser( description='chainer implementation of pix2pix') parser.add_argument('--batchsize', '-b', type=int, default=1, help='Number of images in each mini-batch') parser.add_argument('--epoch', '-e', type=int, default=200, help='Number of sweeps over the dataset to train') parser.add_argument('--device', '-d', type=str, default='-1', help='Device specifier. Either ChainerX device ' 'specifier or an integer. If non-negative integer, ' 'CuPy arrays with specified device id are used. If ' 'negative integer, NumPy arrays are used') parser.add_argument('--dataset', '-i', default='./facade/base', help='Directory of image files.') parser.add_argument('--out', '-o', default='result', help='Directory to output the result') parser.add_argument('--resume', '-r', type=str, help='Resume the training from snapshot') parser.add_argument('--seed', type=int, default=0, help='Random seed') parser.add_argument('--snapshot_interval', type=int, default=1000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=100, help='Interval of displaying log to console') group = parser.add_argument_group('deprecated arguments') group.add_argument('--gpu', '-g', dest='device', type=int, nargs='?', const=0, help='GPU ID (negative value indicates CPU)') args = parser.parse_args() if chainer.get_dtype() == numpy.float16: warnings.warn( 'This example may cause NaN in FP16 mode.', RuntimeWarning) device = chainer.get_device(args.device) if device.xp is chainerx: sys.stderr.write('This example does not support ChainerX devices.\n') sys.exit(1) print('Device: {}'.format(device)) print('# Minibatch-size: {}'.format(args.batchsize)) print('# epoch: {}'.format(args.epoch)) print('') device.use() # Set up a neural network to train enc = Encoder(in_ch=12) dec = Decoder(out_ch=3) dis = Discriminator(in_ch=12, out_ch=3) enc.to_device(device) dec.to_device(device) dis.to_device(device) # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec') return optimizer opt_enc = make_optimizer(enc) opt_dec = make_optimizer(dec) opt_dis = make_optimizer(dis) train_d = FacadeDataset(args.dataset, data_range=(1, 300)) test_d = FacadeDataset(args.dataset, data_range=(300, 379)) train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize) test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize) # Set up a trainer updater = FacadeUpdater( models=(enc, dec, dis), iterator={ 'main': train_iter, 'test': test_iter}, optimizer={ 'enc': opt_enc, 'dec': opt_dec, 'dis': opt_dis}, device=device) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) snapshot_interval = (args.snapshot_interval, 'iteration') display_interval = (args.display_interval, 'iteration') trainer.extend(extensions.snapshot( filename='snapshot_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( enc, 'enc_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dec, 'dec_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'enc/loss', 'dec/loss', 'dis/loss', ]), trigger=display_interval) trainer.extend(extensions.ProgressBar(update_interval=10)) trainer.extend( out_image( updater, enc, dec, 5, 5, args.seed, args.out), trigger=snapshot_interval) if args.resume is not None: # Resume from a snapshot chainer.serializers.load_npz(args.resume, trainer) # Run the training trainer.run()
import numpy as np from net import Encoder, ColorNet, SetmentationNet if __name__ == "__main__": encoder = Encoder() color_decoder = ColorNet(args.feature_dim) segmentator = SegmentationNet(args.feature_dim) criterion = nn.MSELoss() encoder_optim = torch.optim.Adam(encoder.parameters(), lr=args.encoder_lr, weight_decay=1e-5) color_decoder_optim = torch.optim.Adam(color_decoder.parameters(), lr=args.decoder_lr, weight_decay=1e-5) segment_optim = torch.optim.Adam(segmentator.parameters(), lr=args.segmentator_lr, weight_decay=1e-5) data_loader = [] for n, batch in enumerate(data_loader): feature = encoder(batch)
def main(): parser = argparse.ArgumentParser( description='chainer implementation of pix2pix') parser.add_argument('--batchsize', '-b', type=int, default=1, help='Number of images in each mini-batch') parser.add_argument('--epoch', '-e', type=int, default=500, help='Number of sweeps over the dataset to train') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--out', '-o', default='result_dehighlight', help='Directory to output the result') parser.add_argument('--resume', '-r', default='', help='Resume the training from snapshot') parser.add_argument('--seed', type=int, default=0, help='Random seed') parser.add_argument('--snapshot_interval', type=int, default=10000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=10000, help='Interval of displaying log to console') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) print('# Minibatch-size: {}'.format(args.batchsize)) print('# epoch: {}'.format(args.epoch)) print('') # Set up a neural network to train enc = Encoder(in_ch=3) dec = Decoder(out_ch=1) dis = Discriminator(in_ch=3, out_ch=1) gen = Generator(in_ch=3, out_ch=1) serializers.load_npz("depro.npz", depro) gen.encoder = enc gen.decoder = dec if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current enc.to_gpu() # Copy the model to the GPU dec.to_gpu() dis.to_gpu() gen.to_gpu() depro.disable_update() # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec') return optimizer opt_enc = make_optimizer(enc) opt_dec = make_optimizer(dec) opt_dis = make_optimizer(dis) #will fix train_d = NyuDataset("E:/nyu_depth_v2_labeled.mat", startnum=0, endnum=1000) test_d = NyuDataset("E:/nyu_depth_v2_labeled.mat", startnum=1000, endnum=1449) #train_iter = chainer.iterators.MultiprocessIterator(train_d, args.batchsize, n_processes=4) #test_iter = chainer.iterators.MultiprocessIterator(test_d, args.batchsize, n_processes=4) train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize) test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize) test_iter2 = chainer.iterators.SerialIterator(test_d, args.batchsize, repeat=False, shuffle=False) # Set up a trainer updater = PicUpdater(models=(enc, dec, dis), iterator={ 'main': train_iter, 'test': test_iter }, optimizer={ 'enc': opt_enc, 'dec': opt_dec, 'dis': opt_dis }, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) snapshot_interval = (args.snapshot_interval, 'iteration') display_interval = (args.display_interval, 'iteration') trainer.extend( extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( enc, 'enc_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dec, 'dec_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'enc/loss_enc', 'dec/loss_dec', 'dis/loss_dis', "validation/main/loss" ]), trigger=display_interval) trainer.extend(extensions.Evaluator(test_iter2, gen, device=args.gpu)) trainer.extend( extensions.PlotReport(['dec/loss_dec'], x_key='iteration', file_name='dec_loss.png', trigger=display_interval)) trainer.extend( extensions.PlotReport(['dis/loss_dis'], x_key='iteration', file_name='dis_loss.png', trigger=display_interval)) trainer.extend( extensions.PlotReport(["validation/main/loss"], x_key='iteration', file_name='gen_loss.png', trigger=display_interval)) trainer.extend(extensions.ProgressBar(update_interval=10)) trainer.extend(out_image(updater, depro, enc, dec, 3, 3, args.seed, args.out), trigger=snapshot_interval) if args.resume: # Resume from a snapshot chainer.serializers.load_npz(args.resume, trainer) # Run the training trainer.run()
def main(): parser = argparse.ArgumentParser(description='chainer implementation of pix2pix') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--img', '-i', help='Input image') parser.add_argument('--out', '-o', default='result_dehighlight', help='Directory to output the result') args = parser.parse_args() ENC_W = os.path.join(args.out, "enc_iter_2500000.npz") #DEC_W = "trained_model/dec_iter_176000.npz" # to avoid GitHub 100M limit, one .npz files are divided into two zip files. DEC_Ws = [os.path.join(args.out, "dec_iter_2500000.npz")] #shutil.copy("net.py", args.out) # Set up a neural network to train enc = Encoder(in_ch=3) dec = Decoder(out_ch=3) chainer.serializers.load_npz(ENC_W, enc) # to avoid GitHub 100M limit, 1 .npz file is devided into 2 files for npzfile in DEC_Ws: with np.load(npzfile) as f: d = NpzDeserializer(f, strict=False) d.load(dec) if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current enc.to_gpu() # Copy the model to the GPU dec.to_gpu() inimg = loadimg(args.img) ch, h, w = inimg.shape # add paddings so that input array has the size of mutiples of 256. in_ary = np.zeros((ch,math.ceil(h/256)*256, math.ceil(w/256)*256), dtype="f") in_ary[:,0:h,0:w] = inimg x_in = in_ary[np.newaxis,:] # to fit into the minibatch shape print(x_in.shape) # x_in as an input image x_in = chainer.Variable(x_in) if args.gpu >= 0: x_in.to_gpu() st = time.time() for i in tqdm(range(10)): z = enc(x_in) x_out = dec(z) ts = (time.time() - st)/10 print("mean estimation time:{:.2f}".format(ts)) with open(os.path.join(args.out, "time.txt"), "a") as f: f.write("gpu:{}, time:{:.4f}, FPS:{:.4f}\n".format(args.gpu, ts, 1/ts)) if args.gpu >= 0: out_ary = x_out.data.get()[0] else: out_ary = x_out.data[0] #img_show = np.zeros((inimg.shape[0], inimg.shape[1], inimg.shape[2]*2)) #img_show[:,:,:inimg.shape[2]] = inimg #img_show[:,:outimg.shape[1],inimg.shape[2]:inimg.shape[2]+outimg.shape[2]] = outimg outimg = out_ary[:,0:h,0:w] # trim paddings img_show = np.concatenate((inimg, outimg), axis=2) bgrpic = to_bgr(img_show).copy() cv2.putText(bgrpic,"input",(3,15),cv2.FONT_HERSHEY_DUPLEX, 0.5,(255,0,0)) cv2.putText(bgrpic,"output",(w+3,15),cv2.FONT_HERSHEY_DUPLEX, 0.5,(255,0,0)) cv2.imshow("result", bgrpic) cv2.waitKey(0) cv2.destroyAllWindows()
def main(): parser = argparse.ArgumentParser( description='chainer implementation of pix2pix') parser.add_argument('--batchsize', '-b', type=int, default=1, help='Number of images in each mini-batch') parser.add_argument('--epoch', '-e', type=int, default=200, help='Number of sweeps over the dataset to train') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--dataset', '-i', default='./facade/base', help='Directory of image files.') parser.add_argument('--out', '-o', default='result', help='Directory to output the result') parser.add_argument('--resume', '-r', default='', help='Resume the training from snapshot') parser.add_argument('--seed', type=int, default=0, help='Random seed') parser.add_argument('--snapshot_interval', type=int, default=1000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=100, help='Interval of displaying log to console') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) print('# Minibatch-size: {}'.format(args.batchsize)) print('# epoch: {}'.format(args.epoch)) print('') # Set up a neural network to train enc = Encoder(in_ch=12) dec = Decoder(out_ch=3) dis = Discriminator(in_ch=12, out_ch=3) if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current enc.to_gpu() # Copy the model to the GPU dec.to_gpu() dis.to_gpu() # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec') return optimizer opt_enc = make_optimizer(enc) opt_dec = make_optimizer(dec) opt_dis = make_optimizer(dis) train_d = FacadeDataset(args.dataset, data_range=(1, 300)) test_d = FacadeDataset(args.dataset, data_range=(300, 379)) #train_iter = chainer.iterators.MultiprocessIterator(train_d, args.batchsize, n_processes=4) #test_iter = chainer.iterators.MultiprocessIterator(test_d, args.batchsize, n_processes=4) train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize) test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize) # Set up a trainer updater = FacadeUpdater(models=(enc, dec, dis), iterator={ 'main': train_iter, 'test': test_iter }, optimizer={ 'enc': opt_enc, 'dec': opt_dec, 'dis': opt_dis }, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) snapshot_interval = (args.snapshot_interval, 'iteration') display_interval = (args.display_interval, 'iteration') trainer.extend( extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( enc, 'enc_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dec, 'dec_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'enc/loss', 'dec/loss', 'dis/loss', ]), trigger=display_interval) trainer.extend(extensions.ProgressBar(update_interval=10)) trainer.extend(out_image(updater, enc, dec, 5, 5, args.seed, args.out), trigger=snapshot_interval) if args.resume: # Resume from a snapshot chainer.serializers.load_npz(args.resume, trainer) # Run the training trainer.run()
def main(): args = arguments() print(args) if args.imgtype == "dcm": from dataset_dicom import Dataset as Dataset else: from dataset_jpg import DatasetOutMem as Dataset # CUDA if not chainer.cuda.available: print("CUDA required") exit() if len(args.gpu) == 1 and args.gpu[0] >= 0: chainer.cuda.get_device_from_id(args.gpu[0]).use() # Enable autotuner of cuDNN chainer.config.autotune = True chainer.config.dtype = dtypes[args.dtype] chainer.print_runtime_info() # Turn off type check # chainer.config.type_check = False # print('Chainer version: ', chainer.__version__) # print('GPU availability:', chainer.cuda.available) # print('cuDNN availablility:', chainer.cuda.cudnn_enabled) ## dataset iterator print("Setting up data iterators...") train_A_dataset = Dataset(path=os.path.join(args.root, 'trainA'), args=args, base=args.HU_baseA, rang=args.HU_rangeA, random=args.random_translate) train_B_dataset = Dataset(path=os.path.join(args.root, 'trainB'), args=args, base=args.HU_baseB, rang=args.HU_rangeB, random=args.random_translate) test_A_dataset = Dataset(path=os.path.join(args.root, 'testA'), args=args, base=args.HU_baseA, rang=args.HU_rangeA, random=0) test_B_dataset = Dataset(path=os.path.join(args.root, 'testB'), args=args, base=args.HU_baseB, rang=args.HU_rangeB, random=0) args.ch = train_A_dataset.ch args.out_ch = train_B_dataset.ch print("channels in A {}, channels in B {}".format(args.ch, args.out_ch)) # test_A_iter = chainer.iterators.SerialIterator(test_A_dataset, args.nvis_A, shuffle=False) # test_B_iter = chainer.iterators.SerialIterator(test_B_dataset, args.nvis_B, shuffle=False) test_A_iter = chainer.iterators.MultithreadIterator(test_A_dataset, args.nvis_A, shuffle=False, n_threads=3) test_B_iter = chainer.iterators.MultithreadIterator(test_B_dataset, args.nvis_B, shuffle=False, n_threads=3) train_A_iter = chainer.iterators.MultithreadIterator(train_A_dataset, args.batch_size, n_threads=3) train_B_iter = chainer.iterators.MultithreadIterator(train_B_dataset, args.batch_size, n_threads=3) # setup models enc_x = Encoder(args) enc_y = enc_x if args.single_encoder else Encoder(args) dec_x = Decoder(args) dec_y = Decoder(args) dis_x = Discriminator(args) dis_y = Discriminator(args) dis_z = Discriminator( args) if args.lambda_dis_z > 0 else chainer.links.Linear(1, 1) models = { 'enc_x': enc_x, 'dec_x': dec_x, 'enc_y': enc_y, 'dec_y': dec_y, 'dis_x': dis_x, 'dis_y': dis_y, 'dis_z': dis_z } ## load learnt models if args.load_models: for e in models: m = args.load_models.replace('enc_x', e) try: serializers.load_npz(m, models[e]) print('model loaded: {}'.format(m)) except: print("couldn't load {}".format(m)) pass # select GPU if len(args.gpu) == 1: for e in models: models[e].to_gpu() print('using gpu {}, cuDNN {}'.format(args.gpu, chainer.cuda.cudnn_enabled)) else: print("mandatory GPU use: currently only a single GPU can be used") exit() # Setup optimisers def make_optimizer(model, lr, opttype='Adam'): # eps = 1e-5 if args.dtype==np.float16 else 1e-8 optimizer = optim[opttype](lr) #from profiled_optimizer import create_marked_profile_optimizer # optimizer = create_marked_profile_optimizer(optim[opttype](lr), sync=True, sync_level=2) optimizer.setup(model) if args.weight_decay > 0: if opttype in ['Adam', 'AdaBound', 'Eve']: optimizer.weight_decay_rate = args.weight_decay else: if args.weight_decay_norm == 'l2': optimizer.add_hook( chainer.optimizer.WeightDecay(args.weight_decay)) else: optimizer.add_hook( chainer.optimizer_hooks.Lasso(args.weight_decay)) return optimizer opt_enc_x = make_optimizer(enc_x, args.learning_rate_g, args.optimizer) opt_dec_x = make_optimizer(dec_x, args.learning_rate_g, args.optimizer) opt_enc_y = make_optimizer(enc_y, args.learning_rate_g, args.optimizer) opt_dec_y = make_optimizer(dec_y, args.learning_rate_g, args.optimizer) opt_x = make_optimizer(dis_x, args.learning_rate_d, args.optimizer) opt_y = make_optimizer(dis_y, args.learning_rate_d, args.optimizer) opt_z = make_optimizer(dis_z, args.learning_rate_d, args.optimizer) optimizers = { 'opt_enc_x': opt_enc_x, 'opt_dec_x': opt_dec_x, 'opt_enc_y': opt_enc_y, 'opt_dec_y': opt_dec_y, 'opt_x': opt_x, 'opt_y': opt_y, 'opt_z': opt_z } if args.load_optimizer: for e in optimizers: try: m = args.load_models.replace('enc_x', e) serializers.load_npz(m, optimizers[e]) print('optimiser loaded: {}'.format(m)) except: print("couldn't load {}".format(m)) pass # Set up an updater: TODO: multi gpu updater print("Preparing updater...") updater = Updater(models=(enc_x, dec_x, enc_y, dec_y, dis_x, dis_y, dis_z), iterator={ 'main': train_A_iter, 'train_B': train_B_iter, }, optimizer=optimizers, converter=convert.ConcatWithAsyncTransfer(), device=args.gpu[0], params={'args': args}) if args.snapinterval < 0: args.snapinterval = args.lrdecay_start + args.lrdecay_period log_interval = (200, 'iteration') model_save_interval = (args.snapinterval, 'epoch') plot_interval = (500, 'iteration') # Set up a trainer print("Preparing trainer...") if args.iteration: stop_trigger = (args.iteration, 'iteration') else: stop_trigger = (args.lrdecay_start + args.lrdecay_period, 'epoch') trainer = training.Trainer(updater, stop_trigger, out=args.out) for e in models: trainer.extend(extensions.snapshot_object(models[e], e + '{.updater.epoch}.npz'), trigger=model_save_interval) # trainer.extend(extensions.ParameterStatistics(models[e])) ## very slow for e in optimizers: trainer.extend(extensions.snapshot_object(optimizers[e], e + '{.updater.epoch}.npz'), trigger=model_save_interval) log_keys = ['epoch', 'iteration', 'lr'] log_keys_cycle = [ 'opt_enc_x/loss_cycle', 'opt_enc_y/loss_cycle', 'opt_dec_x/loss_cycle', 'opt_dec_y/loss_cycle', 'myval/cycle_x_l1', 'myval/cycle_y_l1' ] log_keys_adv = [ 'opt_enc_y/loss_adv', 'opt_dec_y/loss_adv', 'opt_enc_x/loss_adv', 'opt_dec_x/loss_adv' ] log_keys_d = [] if args.lambda_reg > 0: log_keys.extend(['opt_enc_x/loss_reg', 'opt_enc_y/loss_reg']) if args.lambda_tv > 0: log_keys.extend(['opt_dec_y/loss_tv']) if args.lambda_air > 0: log_keys.extend(['opt_dec_x/loss_air', 'opt_dec_y/loss_air']) if args.lambda_grad > 0: log_keys.extend(['opt_dec_x/loss_grad', 'opt_dec_y/loss_grad']) if args.lambda_identity_x > 0: # perceptual log_keys.extend(['opt_enc_x/loss_id', 'opt_enc_y/loss_id']) if args.lambda_domain > 0: log_keys_cycle.extend(['opt_dec_x/loss_dom', 'opt_dec_y/loss_dom']) if args.dis_reg_weighting > 0: log_keys_d.extend( ['opt_x/loss_reg', 'opt_y/loss_reg', 'opt_z/loss_reg']) if args.dis_wgan: log_keys_d.extend([ 'opt_x/loss_dis', 'opt_x/loss_gp', 'opt_y/loss_dis', 'opt_y/loss_gp' ]) if args.lambda_dis_z > 0: log_keys_d.extend(['opt_z/loss_dis', 'opt_z/loss_gp']) else: log_keys_d.extend([ 'opt_x/loss_real', 'opt_x/loss_fake', 'opt_y/loss_real', 'opt_y/loss_fake' ]) if args.lambda_dis_z > 0: log_keys_d.extend(['opt_z/loss_x', 'opt_z/loss_y']) log_keys_all = log_keys + log_keys_d + log_keys_adv + log_keys_cycle trainer.extend( extensions.LogReport(keys=log_keys_all, trigger=log_interval)) trainer.extend(extensions.PrintReport(log_keys_all), trigger=log_interval) trainer.extend(extensions.ProgressBar(update_interval=20)) trainer.extend(extensions.observe_lr(optimizer_name='opt_enc_x'), trigger=log_interval) # learning rate scheduling decay_start_iter = len(train_A_dataset) * args.lrdecay_start decay_end_iter = len(train_A_dataset) * (args.lrdecay_start + args.lrdecay_period) for e in [opt_enc_x, opt_enc_y, opt_dec_x, opt_dec_y]: trainer.extend( extensions.LinearShift('alpha', (args.learning_rate_g, 0), (decay_start_iter, decay_end_iter), optimizer=e)) for e in [opt_x, opt_y, opt_z]: trainer.extend( extensions.LinearShift('alpha', (args.learning_rate_d, 0), (decay_start_iter, decay_end_iter), optimizer=e)) ## dump graph if args.lambda_Az > 0: trainer.extend( extensions.dump_graph('opt_enc_x/loss_cycle', out_name='gen.dot')) if args.lambda_dis_x > 0: if args.dis_wgan: trainer.extend( extensions.dump_graph('opt_x/loss_dis', out_name='dis.dot')) else: trainer.extend( extensions.dump_graph('opt_x/loss_fake', out_name='dis.dot')) # ChainerUI trainer.extend(CommandsExtension()) if extensions.PlotReport.available(): trainer.extend( extensions.PlotReport(log_keys[3:], 'iteration', trigger=plot_interval, file_name='loss.png', postprocess=plot_log)) trainer.extend( extensions.PlotReport(log_keys_d, 'iteration', trigger=plot_interval, file_name='loss_d.png')) trainer.extend( extensions.PlotReport(log_keys_adv, 'iteration', trigger=plot_interval, file_name='loss_adv.png')) trainer.extend( extensions.PlotReport(log_keys_cycle, 'iteration', trigger=plot_interval, file_name='loss_cyc.png', postprocess=plot_log)) ## visualisation vis_folder = os.path.join(args.out, "vis") os.makedirs(vis_folder, exist_ok=True) if not args.vis_freq: args.vis_freq = len(train_A_dataset) // 2 s = [k for k in range(args.num_slices) ] if args.num_slices > 0 and args.imgtype == "dcm" else None trainer.extend(VisEvaluator({ "testA": test_A_iter, "testB": test_B_iter }, { "enc_x": enc_x, "enc_y": enc_y, "dec_x": dec_x, "dec_y": dec_y }, params={ 'vis_out': vis_folder, 'slice': s, 'args': args }, device=args.gpu[0]), trigger=(args.vis_freq, 'iteration')) ## output filenames of training dataset with open(os.path.join(args.out, 'trainA.txt'), 'w') as output: for f in train_A_dataset.names: output.writelines("\n".join(f)) output.writelines("\n") with open(os.path.join(args.out, 'trainB.txt'), 'w') as output: for f in train_B_dataset.names: output.writelines("\n".join(f)) output.writelines("\n") # archive the scripts rundir = os.path.dirname(os.path.realpath(__file__)) import zipfile with zipfile.ZipFile(os.path.join(args.out, 'script.zip'), 'w', compression=zipfile.ZIP_DEFLATED) as new_zip: for f in [ 'train.py', 'net.py', 'updater.py', 'consts.py', 'losses.py', 'arguments.py', 'convert.py' ]: new_zip.write(os.path.join(rundir, f), arcname=f) # Run the training print("\nresults are saved under: ", args.out) save_args(args, args.out) with open(os.path.join(args.out, "args.txt"), 'w') as fh: fh.write(" ".join(sys.argv)) trainer.run()