def __init__(self): chainer.print_runtime_info() # net としてインスタンス化 self.n_input = 7 self.n_hidden_1 = 5 self.n_hidden_2 = 5 self.n_output = 2 self.net = Sequential( L.Linear(n_input, n_hidden_1), F.relu, L.Linear(n_hidden_1, n_hidden_1), F.relu, L.Linear(n_hidden_1, n_hidden_1), F.relu, L.Linear(n_hidden_1, n_hidden_1), F.relu, L.Linear(n_hidden_1, n_hidden_1), F.relu, L.Linear(n_hidden_1, n_hidden_1), F.relu, L.Linear(n_hidden_1, n_hidden_2), F.relu, L.Linear(n_hidden_2, n_hidden_2), F.relu, L.Linear(n_hidden_2, n_hidden_2), F.relu, L.Linear(n_hidden_2, n_hidden_2), F.relu, L.Linear(n_hidden_2, n_hidden_2), F.relu, L.Linear(n_hidden_2, n_output) )
def main(): args = arguments() outdir = os.path.join(args.out, dt.now().strftime('%m%d_%H%M') + "_cgan") # chainer.config.type_check = False chainer.config.autotune = True chainer.config.dtype = dtypes[args.dtype] chainer.print_runtime_info() #print('Chainer version: ', chainer.__version__) #print('GPU availability:', chainer.cuda.available) #print('cuDNN availability:', chainer.cuda.cudnn_enabled) if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() ## dataset preparation train_d = Dataset(args.train, args.root, args.from_col, args.to_col, clipA=args.clipA, clipB=args.clipB, class_num=args.class_num, crop=(args.crop_height, args.crop_width), imgtype=args.imgtype, random=args.random_translate, grey=args.grey, BtoA=args.btoa) test_d = Dataset(args.val, args.root, args.from_col, args.to_col, clipA=args.clipA, clipB=args.clipB, class_num=args.class_num, crop=(args.crop_height, args.crop_width), imgtype=args.imgtype, random=args.random_translate, grey=args.grey, BtoA=args.btoa) args.crop_height, args.crop_width = train_d.crop if (len(train_d) == 0): print("No images found!") exit() # setup training/validation data iterators train_iter = chainer.iterators.SerialIterator(train_d, args.batch_size) test_iter = chainer.iterators.SerialIterator(test_d, args.nvis, shuffle=False) test_iter_gt = chainer.iterators.SerialIterator( train_d, args.nvis, shuffle=False) ## same as training data; used for validation args.ch = len(train_d[0][0]) args.out_ch = len(train_d[0][1]) print("Input channels {}, Output channels {}".format(args.ch, args.out_ch)) if (len(train_d) * len(test_d) == 0): print("No images found!") exit() ## Set up models # shared pretrained layer if (args.gen_pretrained_encoder and args.gen_pretrained_lr_ratio == 0): if "resnet" in args.gen_pretrained_encoder: pretrained = L.ResNet50Layers() print("Pretrained ResNet model loaded.") else: pretrained = L.VGG16Layers() print("Pretrained VGG model loaded.") if args.gpu >= 0: pretrained.to_gpu() enc_x = net.Encoder(args, pretrained) else: enc_x = net.Encoder(args) # gen = net.Generator(args) dec_y = net.Decoder(args) if args.lambda_dis > 0: dis = net.Discriminator(args) models = {'enc_x': enc_x, 'dec_y': dec_y, 'dis': dis} else: dis = L.Linear(1, 1) models = {'enc_x': enc_x, 'dec_y': dec_y} ## load learnt models optimiser_files = [] if args.model_gen: serializers.load_npz(args.model_gen, enc_x) serializers.load_npz(args.model_gen.replace('enc_x', 'dec_y'), dec_y) print('model loaded: {}, {}'.format( args.model_gen, args.model_gen.replace('enc_x', 'dec_y'))) optimiser_files.append(args.model_gen.replace('enc_x', 'opt_enc_x')) optimiser_files.append(args.model_gen.replace('enc_x', 'opt_dec_y')) if args.model_dis: serializers.load_npz(args.model_dis, dis) print('model loaded: {}'.format(args.model_dis)) optimiser_files.append(args.model_dis.replace('dis', 'opt_dis')) ## send models to GPU if args.gpu >= 0: enc_x.to_gpu() dec_y.to_gpu() dis.to_gpu() # Setup optimisers def make_optimizer(model, lr, opttype='Adam', pretrained_lr_ratio=1.0): # eps = 1e-5 if args.dtype==np.float16 else 1e-8 optimizer = optim[opttype](lr) 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_gen, args.optimizer) opt_dec_y = make_optimizer(dec_y, args.learning_rate_gen, args.optimizer) opt_dis = make_optimizer(dis, args.learning_rate_dis, args.optimizer) optimizers = {'enc_x': opt_enc_x, 'dec_y': opt_dec_y, 'dis': opt_dis} ## resume optimisers from file if args.load_optimizer: for (m, e) in zip(optimiser_files, optimizers): if m: try: serializers.load_npz(m, optimizers[e]) print('optimiser loaded: {}'.format(m)) except: print("couldn't load {}".format(m)) pass # finetuning if args.gen_pretrained_encoder: if args.gen_pretrained_lr_ratio == 0: enc_x.base.disable_update() else: for func_name in enc_x.encoder.base._children: for param in enc_x.encoder.base[func_name].params(): param.update_rule.hyperparam.eta *= args.gen_pretrained_lr_ratio # Set up trainer updater = Updater( models=(enc_x, dec_y, dis), iterator={'main': train_iter}, optimizer=optimizers, # converter=convert.ConcatWithAsyncTransfer(), params={'args': args}, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=outdir) ## save learnt results at a specified interval or at the end of training if args.snapinterval < 0: args.snapinterval = args.epoch snapshot_interval = (args.snapinterval, 'epoch') display_interval = (args.display_interval, 'iteration') for e in models: trainer.extend(extensions.snapshot_object(models[e], e + '{.updater.epoch}.npz'), trigger=snapshot_interval) if args.parameter_statistics: trainer.extend(extensions.ParameterStatistics( models[e])) ## very slow for e in optimizers: trainer.extend(extensions.snapshot_object( optimizers[e], 'opt_' + e + '{.updater.epoch}.npz'), trigger=snapshot_interval) ## plot NN graph if args.lambda_rec_l1 > 0: trainer.extend( extensions.dump_graph('dec_y/loss_L1', out_name='enc.dot')) elif args.lambda_rec_l2 > 0: trainer.extend( extensions.dump_graph('dec_y/loss_L2', out_name='gen.dot')) elif args.lambda_rec_ce > 0: trainer.extend( extensions.dump_graph('dec_y/loss_CE', out_name='gen.dot')) if args.lambda_dis > 0: trainer.extend( extensions.dump_graph('dis/loss_real', out_name='dis.dot')) ## log outputs log_keys = ['epoch', 'iteration', 'lr'] log_keys_gen = ['myval/loss_L1', 'myval/loss_L2'] log_keys_dis = [] if args.lambda_rec_l1 > 0: log_keys_gen.append('dec_y/loss_L1') if args.lambda_rec_l2 > 0: log_keys_gen.append('dec_y/loss_L2') if args.lambda_rec_ce > 0: log_keys_gen.extend(['dec_y/loss_CE', 'myval/loss_CE']) if args.lambda_reg > 0: log_keys.extend(['enc_x/loss_reg']) if args.lambda_tv > 0: log_keys_gen.append('dec_y/loss_tv') if args.lambda_dis > 0: log_keys_dis.extend( ['dec_y/loss_dis', 'dis/loss_real', 'dis/loss_fake']) if args.lambda_mispair > 0: log_keys_dis.append('dis/loss_mispair') if args.dis_wgan: log_keys_dis.extend(['dis/loss_gp']) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport(log_keys + log_keys_gen + log_keys_dis), trigger=display_interval) if extensions.PlotReport.available(): # trainer.extend(extensions.PlotReport(['lr'], 'iteration',trigger=display_interval, file_name='lr.png')) trainer.extend( extensions.PlotReport(log_keys_gen, 'iteration', trigger=display_interval, file_name='loss_gen.png', postprocess=plot_log)) trainer.extend( extensions.PlotReport(log_keys_dis, 'iteration', trigger=display_interval, file_name='loss_dis.png')) trainer.extend(extensions.ProgressBar(update_interval=10)) # learning rate scheduling trainer.extend(extensions.observe_lr(optimizer_name='enc_x'), trigger=display_interval) if args.optimizer in ['Adam', 'AdaBound', 'Eve']: lr_target = 'eta' else: lr_target = 'lr' if args.lr_drop > 0: ## cosine annealing for e in [opt_enc_x, opt_dec_y, opt_dis]: trainer.extend(CosineShift(lr_target, args.epoch // args.lr_drop, optimizer=e), trigger=(1, 'epoch')) else: for e in [opt_enc_x, opt_dec_y, opt_dis]: #trainer.extend(extensions.LinearShift('eta', (1.0,0.0), (decay_start_iter,decay_end_iter), optimizer=e)) trainer.extend(extensions.ExponentialShift('lr', 0.33, optimizer=e), trigger=(args.epoch // args.lr_drop, 'epoch')) # evaluation vis_folder = os.path.join(outdir, "vis") os.makedirs(vis_folder, exist_ok=True) if not args.vis_freq: args.vis_freq = max(len(train_d) // 2, 50) trainer.extend(VisEvaluator({ "test": test_iter, "train": test_iter_gt }, { "enc_x": enc_x, "dec_y": dec_y }, params={ 'vis_out': vis_folder, 'args': args }, device=args.gpu), trigger=(args.vis_freq, 'iteration')) # ChainerUI: removed until ChainerUI updates to be compatible with Chainer 6.0 trainer.extend(CommandsExtension()) # Run the training print("\nresults are saved under: ", outdir) save_args(args, outdir) trainer.run()
from chainer import report from chainer import training from chainer import utils from chainer import Variable from chainer import datasets, iterators, optimizers, serializers from chainer import Link, Chain, ChainList, cuda import chainer.functions as F import chainer.links as L from chainer.training import extensions from PIL import Image import matplotlib.pyplot as plt # chainerのバージョンの確認 print('GPU availability:', chainer.cuda.available) print('cuDNN availablility:', chainer.cuda.cudnn_enabled) chainer.print_runtime_info() # CPU,GPUの確認 !cat /proc/cpuinfo !cat /proc/driver/nvidia/gpus/0000:00:04.0/information # クラス数 class_num = 10 # 画像の大きさ XSIZE = 32 YSIZE = 32 # 学習データ数 train_num = 200
def main(): import chainermn chainer.global_config.autotune = True parser = argparse.ArgumentParser(description='ChainerMN example: Train MQAP using 3DCNN') parser.add_argument('--communicator', type=str, default='hierarchical', help='Type of communicator') parser.add_argument('--gpu', '-g', action='store_true', help='Use GPU') parser.add_argument('--out', '-o', default='result', help='Directory to output the result') parser.add_argument('--resume', '-r', action='store_true', help='Resume the training from snapshot') parser.add_argument('--weight', '-w', action='store_true', help='Resume only weight') parser.add_argument('--config', '-c', type=int, default=0, help='Number of config') parser.add_argument('--config_file', type=str, default='./data/config.json', help='Config file path') args = parser.parse_args() if args.gpu: if args.communicator == 'naive': print("Error: 'naive' communicator does not support GPU.\n") exit(-1) comm = chainermn.create_communicator(args.communicator, allreduce_grad_dtype='float16') device = comm.intra_rank else: if args.communicator != 'naive': print('Warning: using naive communicator ' 'because only naive supports CPU-only execution') comm = chainermn.create_communicator('naive') device = -1 f = open(args.config_file, 'r') config = json.load(f)['Config'][args.config] args.out = os.path.join(args.out, str(args.config)) if comm.rank == 0: print('==========================================') chainer.print_runtime_info() print('Num process (COMM_WORLD): {}'.format(comm.size)) if args.gpu: print('Using GPUs') print('Using {} communicator'.format(args.communicator)) print('Num epoch: {}'.format(config['epoch'])) print('Batch size: {}'.format(config['batch_size'] * comm.size)) print('Optimizer: {}'.format(config['optimizer'])) print('Learning Rate: {}'.format(config['learning_rate'])) print('Out Directory: {}'.format(args.out)) print('Vertex feature: {}'.format(config['vertex_feature'])) if config['global_mode']: print('Using Global loss') if config['local_mode']: print('Using local loss') print('Local type : {}'.format(config['local_type'])) print('Local label : {}'.format(config['local_label'])) print('==========================================') d = Dataproc(size=comm.size, rank=comm.rank, config=config) if device >= 0: chainer.cuda.get_device(device).use() # sub_comm = comm.split(comm.rank // comm.intra_size, comm.rank) if config['local_type'] == 'Regression': local_loss_func = F.mean_squared_error else: local_loss_func = F.sigmoid_cross_entropy global_loss_func = F.mean_squared_error model = build_model(config=config, comm=comm) model = Classifier(predictor=model, local_loss_func=local_loss_func, global_loss_func=global_loss_func, config=config) if device >= 0: model.to_gpu() train, test = d.get_dataset(key='train'), d.get_dataset(key='test') train_iter = I.SerialIterator(dataset=train, batch_size=config['batch_size'], repeat=True, shuffle=True) test_iter = I.SerialIterator(dataset=test, batch_size=config['batch_size'], repeat=False, shuffle=False) # train_iter = I.MultiprocessIterator(dataset=train, batch_size=args.batch, repeat=True, shuffle=True, n_processes=10) # test_iter = I.MultiprocessIterator(dataset=test, batch_size=args.batch, repeat=False, shuffle=True, n_processes=10) if config['optimizer'] == 'Adam': optimizer = chainer.optimizers.Adam(alpha=config['learning_rate'], weight_decay_rate=config['weight_decay_rate'], amsgrad=True) optimizer = chainermn.create_multi_node_optimizer(optimizer, comm, double_buffering=False) elif config['optimizer'] == 'MomentumSGD': optimizer = chainer.optimizers.MomentumSGD(lr=config['learning_rate']) optimizer = chainermn.create_multi_node_optimizer(optimizer, comm, double_buffering=False) elif config['optimizer'] == 'SMORMS3': optimizer = chainer.optimizers.SMORMS3(lr=config['learning_rate']) optimizer = chainermn.create_multi_node_optimizer(optimizer, comm, double_buffering=False) elif config['optimizer'] == 'Eve': from my_optimizer.eve import Eve, create_multi_node_optimizer optimizer = Eve(alpha=config['learning_rate']) optimizer = create_multi_node_optimizer(optimizer, comm, double_buffering=False) elif config['optimizer'] == 'Adabound': from my_optimizer.adabound import Adam as Adabound optimizer = Adabound(alpha=config['learning_rate'], adabound=True, amsgrad=True, weight_decay_rate=config['weight_decay_rate']) optimizer = chainermn.create_multi_node_optimizer(optimizer, comm, double_buffering=False) optimizer.setup(model) val_interval = 1, 'epoch' log_interval = 1, 'epoch' updater = training.StandardUpdater(train_iter, optimizer, device=device, converter=d.get_converter()) trainer = training.Trainer(updater, (config['epoch'], 'epoch'), out=args.out) evaluator = GraphEvaluator(iterator=test_iter, target=model.predictor, device=device, converter=d.get_converter(), comm=comm, local_loss_func=local_loss_func, global_loss_func=global_loss_func, name='val', config=config) evaluator = chainermn.create_multi_node_evaluator(evaluator, comm) trainer.extend(evaluator, trigger=val_interval) if comm.rank == 0: trainer.extend(extensions.dump_graph('main/loss')) trainer.extend(extensions.snapshot(), trigger=val_interval) trainer.extend(extensions.LogReport(trigger=log_interval)) trainer.extend(extensions.PlotReport(['main/loss', 'val/main/loss'], 'epoch', file_name='loss.png'), trigger=val_interval) report_list = ['epoch', 'main/loss', 'val/main/loss'] if config['global_mode']: report_list.extend(['main/global_loss', 'val/main/global_loss', 'val/main/global_pearson']) trainer.extend(extensions.PlotReport(['main/global_loss', 'val/main/global_loss'], 'epoch', file_name='global_loss.png'), trigger=val_interval) if config['local_mode']: report_list.extend(['main/local_loss', 'val/main/local_loss', 'val/main/local_mean_pearson']) if config['local_type'] == 'Classification': report_list.append('val/main/local_auc') trainer.extend(extensions.PlotReport(['val/main/local_auc'], 'epoch', file_name='local_auc.png'), trigger=val_interval) else: report_list.append('val/main/local_pearson') report_list.append('elapsed_time') trainer.extend(extensions.PrintReport(report_list), trigger=log_interval) trainer.extend(extensions.ProgressBar(update_interval=10)) if args.resume: snap_list = [p for p in os.listdir(args.out) if 'snapshot' in p] snap_num = np.array([int(re.findall("[+-]?[0-9]+[\.]?[0-9]*[eE]?[+-]?[0-9]*", p)[0]) for p in snap_list]) path = snap_list[np.argmax(snap_num)] path = os.path.join(args.out, path) if args.weight: obj_path = 'updater/model:main/predictor/' chainer.serializers.load_npz(path, model.predictor, obj_path) else: chainer.serializers.load_npz(path, trainer) if comm.rank == 0: protein_name_dict = d.get_protein_name_dict() out_path = Path(args.out) if not out_path.exists(): out_path.mkdir(parents=True, exist_ok=True) np.savez(os.path.join(args.out, 'protein_name'), **protein_name_dict) f = open(os.path.join(args.out, 'config.json'), 'w') json.dump(config, f, ensure_ascii=False, indent=4, sort_keys=True, separators=(',', ': ')) f.close() f = open(os.path.join(args.out, 'args.json'), 'w') json.dump(vars(args), f) f.close() if comm.rank == 0: print('train start!!!') trainer.run()
def main(): args = arguments() # chainer.config.type_check = False chainer.config.autotune = True chainer.config.dtype = dtypes[args.dtype] chainer.print_runtime_info() #print('Chainer version: ', chainer.__version__) #print('GPU availability:', chainer.cuda.available) #print('cuDNN availability:', chainer.cuda.cudnn_enabled) ## dataset preparation if args.imgtype == "dcm": from dataset_dicom import Dataset else: from dataset import Dataset train_d = Dataset(args.train, args.root, args.from_col, args.to_col, crop=(args.crop_height, args.crop_width), random=args.random_translate, grey=args.grey) test_d = Dataset(args.val, args.root, args.from_col, args.to_col, crop=(args.crop_height, args.crop_width), random=args.random_translate, grey=args.grey) # setup training/validation data iterators train_iter = chainer.iterators.SerialIterator(train_d, args.batch_size) test_iter = chainer.iterators.SerialIterator(test_d, args.nvis, shuffle=False) test_iter_gt = chainer.iterators.SerialIterator( train_d, args.nvis, shuffle=False) ## same as training data; used for validation args.ch = len(train_d[0][0]) args.out_ch = len(train_d[0][1]) print("Input channels {}, Output channels {}".format(args.ch, args.out_ch)) ## Set up models gen = net.Generator(args) dis = net.Discriminator(args) ## load learnt models optimiser_files = [] if args.model_gen: serializers.load_npz(args.model_gen, gen) print('model loaded: {}'.format(args.model_gen)) optimiser_files.append(args.model_gen.replace('gen_', 'opt_gen_')) if args.model_dis: serializers.load_npz(args.model_dis, dis) print('model loaded: {}'.format(args.model_dis)) optimiser_files.append(args.model_dis.replace('dis_', 'opt_dis_')) ## send models to GPU if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() gen.to_gpu() dis.to_gpu() # Setup optimisers def make_optimizer(model, lr, opttype='Adam'): # eps = 1e-5 if args.dtype==np.float16 else 1e-8 optimizer = optim[opttype](lr) 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)) optimizer.setup(model) return optimizer opt_gen = make_optimizer(gen, args.learning_rate, args.optimizer) opt_dis = make_optimizer(dis, args.learning_rate, args.optimizer) optimizers = {'opt_g': opt_gen, 'opt_d': opt_dis} ## resume optimisers from file if args.load_optimizer: for (m, e) in zip(optimiser_files, optimizers): if m: try: serializers.load_npz(m, optimizers[e]) print('optimiser loaded: {}'.format(m)) except: print("couldn't load {}".format(m)) pass # Set up trainer updater = pixupdater( models=(gen, dis), iterator={ 'main': train_iter, 'test': test_iter, 'test_gt': test_iter_gt }, optimizer={ 'gen': opt_gen, 'dis': opt_dis }, # converter=convert.ConcatWithAsyncTransfer(), params={'args': args}, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) ## save learnt results at an interval if args.snapinterval < 0: args.snapinterval = args.epoch snapshot_interval = (args.snapinterval, 'epoch') display_interval = (args.display_interval, 'iteration') trainer.extend(extensions.snapshot_object(gen, 'gen_{.updater.epoch}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object(opt_gen, 'opt_gen_{.updater.epoch}.npz'), trigger=snapshot_interval) if args.lambda_dis > 0: trainer.extend(extensions.snapshot_object(dis, 'dis_{.updater.epoch}.npz'), trigger=snapshot_interval) trainer.extend( extensions.dump_graph('dis/loss_real', out_name='dis.dot')) trainer.extend(extensions.snapshot_object( opt_dis, 'opt_dis_{.updater.epoch}.npz'), trigger=snapshot_interval) if args.lambda_rec_l1 > 0: trainer.extend(extensions.dump_graph('gen/loss_L1', out_name='gen.dot')) elif args.lambda_rec_l2 > 0: trainer.extend(extensions.dump_graph('gen/loss_L2', out_name='gen.dot')) ## log outputs log_keys = ['epoch', 'iteration', 'lr'] log_keys_gen = [ 'gen/loss_L1', 'gen/loss_L2', 'gen/loss_dis', 'myval/loss_L2', 'gen/loss_tv' ] log_keys_dis = ['dis/loss_real', 'dis/loss_fake', 'dis/loss_mispair'] trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport(log_keys + log_keys_gen + log_keys_dis), trigger=display_interval) if extensions.PlotReport.available(): trainer.extend( extensions.PlotReport(log_keys_gen, 'iteration', trigger=display_interval, file_name='loss_gen.png')) trainer.extend( extensions.PlotReport(log_keys_dis, 'iteration', trigger=display_interval, file_name='loss_dis.png')) trainer.extend(extensions.ProgressBar(update_interval=10)) trainer.extend(extensions.ParameterStatistics(gen)) # learning rate scheduling if args.optimizer in ['SGD', 'Momentum', 'AdaGrad', 'RMSprop']: trainer.extend(extensions.observe_lr(optimizer_name='gen'), trigger=display_interval) trainer.extend(extensions.ExponentialShift('lr', 0.33, optimizer=opt_gen), trigger=(args.epoch / 5, 'epoch')) trainer.extend(extensions.ExponentialShift('lr', 0.33, optimizer=opt_dis), trigger=(args.epoch / 5, 'epoch')) elif args.optimizer in ['Adam', 'AdaBound', 'Eve']: trainer.extend(extensions.observe_lr(optimizer_name='gen'), trigger=display_interval) trainer.extend(extensions.ExponentialShift("alpha", 0.33, optimizer=opt_gen), trigger=(args.epoch / 5, 'epoch')) trainer.extend(extensions.ExponentialShift("alpha", 0.33, optimizer=opt_dis), trigger=(args.epoch / 5, 'epoch')) # evaluation 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_d) // 2 trainer.extend(VisEvaluator({ "test": test_iter, "train": test_iter_gt }, {"gen": gen}, params={'vis_out': vis_folder}, device=args.gpu), trigger=(args.vis_freq, 'iteration')) # ChainerUI: removed until ChainerUI updates to be compatible with Chainer 6.0 # trainer.extend(CommandsExtension()) # Run the training print("trainer start") trainer.run()
def dump_chainer_info(self): """ returns chainer, CuPy and CuDNN version info """ chainer.print_runtime_info()
def gqn_process(): # load model my_gpu = args.gpu_device if my_gpu < 0: xp = np else: cuda.get_device(args.gpu_device).use() xp = cp hyperparams = HyperParameters() assert hyperparams.load(args.snapshot_directory) model = Model(hyperparams) chainer.serializers.load_hdf5(args.snapshot_file, model) if my_gpu > -1: model.to_gpu() chainer.print_runtime_info() observed_viewpoint, observed_image, offset = data_recv.get() observed_viewpoint = np.expand_dims(np.expand_dims( np.asarray(observed_viewpoint).astype(np.float32), axis=0), axis=0) observed_image = np.expand_dims(np.expand_dims( np.asarray(observed_image).astype(np.float32), axis=0), axis=0) offset = np.asarray(offset) camera_distance = np.mean( np.linalg.norm(observed_viewpoint[:, :, :3], axis=2)) camera_position_z = np.mean(observed_viewpoint[:, :, 1]) observed_image = observed_image.transpose( (0, 1, 4, 2, 3)).astype(np.float32) observed_image = preprocess_images(observed_image) # create representation and generate uncertainty map of environment [1000 viewpoints?] total_frames = 10 representation = model.compute_observation_representation( observed_image, observed_viewpoint) # get predictions highest_var = 0.0 no_of_samples = 20 highest_var_vp = 0 try: for i in range(0, total_frames): horizontal_angle_rad = compute_camera_angle_at_frame( i, total_frames) query_viewpoints = rotate_query_viewpoint(horizontal_angle_rad, camera_distance, camera_position_z, xp) generated_images = xp.squeeze( xp.array( model.generate_images(query_viewpoints, representation, no_of_samples))) var_image = xp.var(generated_images, axis=0) # var_image = chainer.backends.cuda.to_cpu(var_image) # grayscale # r,g,b = var_image # gray_var_image = 0.2989*r+0.5870*g+0.1140*b current_var = xp.mean(var_image) if highest_var == 0: highest_var = current_var highest_var_vp = query_viewpoints[0] elif current_var > highest_var: highest_var = current_var highest_var_vp = query_viewpoints[0] except KeyboardInterrupt: logging.warning('interrupt') # return next viewpoint and unit vector of end effector based on highest uncertainty found in the uncertainty map _x, _y, _z, _, _, _, _ = highest_var_vp _yaw, _pitch = compute_yaw_and_pitch([_x, _y, _z]) next_viewpoint = [_x, _y, _z, _yaw, _pitch] next_viewpoint = [chainer.backends.cuda.to_cpu(x) for x in next_viewpoint] next_viewpoint = [float(x) for x in next_viewpoint] data_send.put(next_viewpoint)
def main(): print("Hello Chainer") chainer.print_runtime_info()
def main(): args = arguments() out = os.path.join(args.out, dt.now().strftime('%m%d_%H%M_AE')) print(args) print(out) save_args(args, out) args.dtype = dtypes[args.dtype] args.dis_activation = activation[args.dis_activation] args.gen_activation = activation[args.gen_activation] args.gen_out_activation = activation[args.gen_out_activation] args.gen_nblock = args.gen_nblock // 2 # to match ordinary cycleGAN if args.imgtype=="dcm": from dataset_dicom import DatasetOutMem as Dataset else: from dataset_jpg import DatasetOutMem as Dataset if not chainer.cuda.available: print("CUDA required") 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 = 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'), baseA=args.HU_base, rangeA=args.HU_range, slice_range=args.slice_range, crop=(args.crop_height,args.crop_width),random=args.random_translate, forceSpacing=0, imgtype=args.imgtype, dtype=args.dtype) train_B_dataset = Dataset( path=os.path.join(args.root, 'trainB'), baseA=args.HU_base, rangeA=args.HU_range, slice_range=args.slice_range, crop=(args.crop_height,args.crop_width), random=args.random_translate, forceSpacing=args.forceSpacing, imgtype=args.imgtype, dtype=args.dtype) test_A_dataset = Dataset( path=os.path.join(args.root, 'testA'), baseA=args.HU_base, rangeA=args.HU_range, slice_range=args.slice_range, crop=(args.crop_height,args.crop_width), random=0, forceSpacing=0, imgtype=args.imgtype, dtype=args.dtype) test_B_dataset = Dataset( path=os.path.join(args.root, 'testB'), baseA=args.HU_base, rangeA=args.HU_range, slice_range=args.slice_range, crop=(args.crop_height,args.crop_width), random=0, forceSpacing=args.forceSpacing, imgtype=args.imgtype, dtype=args.dtype) args.ch = train_A_dataset.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) if args.batch_size > 1: train_A_iter = chainer.iterators.MultiprocessIterator( train_A_dataset, args.batch_size, n_processes=3, shuffle=not args.conditional_discriminator) train_B_iter = chainer.iterators.MultiprocessIterator( train_B_dataset, args.batch_size, n_processes=3, shuffle=not args.conditional_discriminator) else: train_A_iter = chainer.iterators.SerialIterator( train_A_dataset, args.batch_size, shuffle=not args.conditional_discriminator) train_B_iter = chainer.iterators.SerialIterator( train_B_dataset, args.batch_size, shuffle=not args.conditional_discriminator) # setup models enc_x = net.Encoder(args) enc_y = net.Encoder(args) dec_x = net.Decoder(args) dec_y = net.Decoder(args) dis_x = net.Discriminator(args) dis_y = net.Discriminator(args) dis_z = net.Discriminator(args) models = {'enc_x': enc_x, 'enc_y': enc_y, 'dec_x': dec_x, 'dec_y': dec_y, 'dis_x': dis_x, 'dis_y': dis_y, 'dis_z': dis_z} optimiser_files = [] ## 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 optimiser_files.append(m.replace(e,'opt_'+e).replace('dis_','')) # select GPU if len(args.gpu) == 1: for e in models: models[e].to_gpu() print('use 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, alpha=0.0002, beta1=0.5): eps = 1e-5 if args.dtype==np.float16 else 1e-8 optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1, eps=eps) optimizer.setup(model) if args.weight_decay>0: 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, alpha=args.learning_rate_g) opt_dec_x = make_optimizer(dec_x, alpha=args.learning_rate_g) opt_enc_y = make_optimizer(enc_y, alpha=args.learning_rate_g) opt_dec_y = make_optimizer(dec_y, alpha=args.learning_rate_g) opt_x = make_optimizer(dis_x, alpha=args.learning_rate_d) opt_y = make_optimizer(dis_y, alpha=args.learning_rate_d) opt_z = make_optimizer(dis_z, alpha=args.learning_rate_d) 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 (m,e) in zip(optimiser_files,optimizers): if m: try: 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') vis_interval = (args.vis_freq, 'iteration') plot_interval = (500, 'iteration') # Set up a trainer print("Preparing trainer...") trainer = training.Trainer(updater, (args.lrdecay_start + args.lrdecay_period, 'epoch'), out=out) for e in models: trainer.extend(extensions.snapshot_object( models[e], e+'{.updater.epoch}.npz'), trigger=model_save_interval) for e in optimizers: trainer.extend(extensions.snapshot_object( optimizers[e], e+'{.updater.epoch}.npz'), trigger=model_save_interval) log_keys = ['epoch', 'iteration'] 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_y_l1'] log_keys_d = ['opt_x/loss_real','opt_x/loss_fake','opt_y/loss_real','opt_y/loss_fake','opt_z/loss_x','opt_z/loss_y'] 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.extend([ 'opt_dec_y/loss_id']) log_keys.extend([ 'opt_enc_x/loss_reg','opt_enc_y/loss_reg', 'opt_dec_x/loss_air','opt_dec_y/loss_air', 'opt_dec_y/loss_tv']) log_keys_d.extend(['opt_x/loss_gp','opt_y/loss_gp']) 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(CommandsExtension()) ## to dump graph, set -lix 1 --warmup 0 # trainer.extend(extensions.dump_graph('opt_g/loss_id', out_name='gen.dot')) # trainer.extend(extensions.dump_graph('opt_x/loss', out_name='dis.dot')) if extensions.PlotReport.available(): trainer.extend(extensions.PlotReport(log_keys[2:], 'iteration',trigger=plot_interval, file_name='loss.png')) 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')) ## output filenames of training dataset with open(os.path.join(out, 'trainA.txt'),'w') as output: output.writelines("\n".join(train_A_dataset.ids)) with open(os.path.join(out, 'trainB.txt'),'w') as output: output.writelines("\n".join(train_B_dataset.ids)) # archive the scripts rundir = os.path.dirname(os.path.realpath(__file__)) import zipfile with zipfile.ZipFile(os.path.join(out,'script.zip'), 'w', compression=zipfile.ZIP_DEFLATED) as new_zip: for f in ['trainAE.py','net.py','updaterAE.py','consts.py','losses.py','arguments.py','convert.py']: new_zip.write(os.path.join(rundir,f),arcname=f) ## visualisation vis_folder = os.path.join(out, "vis") if not os.path.exists(vis_folder): os.makedirs(vis_folder) # trainer.extend(visualize( (enc_x, enc_y, dec_y), vis_folder, test_A_iter, test_B_iter),trigger=(1, 'epoch')) trainer.extend(VisEvaluator({"main":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, 'single_encoder': args.single_encoder}, device=args.gpu[0]),trigger=vis_interval) # Run the training trainer.run()
def main(): args = arguments() chainer.config.autotune = True chainer.print_runtime_info() print(args) if args.dp: from net_dp import Encoder, Decoder else: from net import Encoder, Decoder if args.gpu >= 0: cuda.get_device(args.gpu).use() xp = cuda.cupy sp = cupyx.scipy.sparse else: print("runs desperately slowly without a GPU!") xp = np sp = scipy.sparse ## Input information ## # InputFile = scanconf.ScanConfig() # InputFile.reconSize = args.crop_width ## setup trainable links decoder = Decoder(args) if args.use_enc: encoder = Encoder(args) else: encoder = L.Linear(1) if args.use_dis: dis = Discriminator(args) else: dis = L.Linear(1) if args.model_dis: serializers.load_npz(args.model_dis, dis) print('discriminator model loaded: {}'.format(args.model_dis)) if args.model_gen: if 'enc' in args.model_gen and not args.decoder_only: serializers.load_npz(args.model_gen, encoder) print('encoder model loaded: {}'.format(args.model_gen)) serializers.load_npz(args.model_gen.replace('enc', 'dec'), decoder) print('decoder model loaded: {}'.format( args.model_gen.replace('enc', 'dec'))) if args.lambda_sd > 0 and args.lr_sd < 0.05: print( "\n\n for usual iterative reconstruction (-ls), --lr_sd should be around 0.1. \n\n" ) if args.latent_dim > 0: init = xp.zeros((args.batchsize, args.latent_dim)).astype(np.float32) elif args.decoder_only: init = xp.zeros((args.batchsize, decoder.latent_c, decoder.latent_h, decoder.latent_w)).astype(np.float32) else: init = xp.zeros((args.batchsize, 1, args.crop_height, args.crop_width)).astype(np.float32) # init = xp.random.uniform(-0.1,0.1,(1,1,args.crop_height,args.crop_width)).astype(np.float32) print("Initial image {} shape {}".format(args.model_image, init.shape)) seed = L.Parameter(init) if args.gpu >= 0: decoder.to_gpu() seed.to_gpu() encoder.to_gpu() dis.to_gpu() # 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', 'Adam_d', 'AdaBound', 'Eve']: optimizer.weight_decay_rate = args.weight_decay else: optimizer.add_hook( chainer.optimizer_hooks.WeightDecay(args.weight_decay)) # optimizer.add_hook(chainer.optimizer_hooks.GradientClipping(100)) return optimizer optimizer_sd = make_optimizer(seed, args.lr_sd, args.optimizer) optimizer_dec = make_optimizer(decoder, args.lr_gen, args.optimizer) optimizer_enc = make_optimizer(encoder, args.lr_gen, args.optimizer) optimizer_dis = make_optimizer(dis, args.lr_dis, args.optimizer_dis) # unify CPU and GPU memory to load big matrices if args.unified_memory_pool and args.crop_height > 256: pool = cp.cuda.MemoryPool(cp.cuda.malloc_managed) cp.cuda.set_allocator(pool.malloc) # projection matrices prMats, conjMats = None, None if args.lambda_sd > 0 or args.lambda_nn > 0: prMat = scipy.sparse.load_npz( os.path.join(args.root, args.projection_matrix)).tocsr(copy=False) # cx = prMat.tocsr() # rows,cols = cx.nonzero() # for i,j in zip(rows,cols): # if cx[i,j] < 1e-5: # cx[i,j] = 0 # prMat = cx.tocoo() # scipy.sparse.save_npz("d:/ml/reconst/pr.npz",prMat) prMats = [ sp.coo_matrix((prMat[np.arange(i, prMat.shape[0], args.osem), :]), dtype=np.float32) for i in range(args.osem) ] prMats = [ chainer.utils.CooMatrix(p.data, p.row, p.col, p.shape) for p in prMats ] print("Projection matrix {} shape {}, thinned {} x {}".format( args.projection_matrix, prMat.shape, prMats[0].shape, len(prMats))) if args.system_matrix: conjMat = scipy.sparse.load_npz( os.path.join(args.root, args.system_matrix)).tocsr(copy=False) conjMats = [ sp.coo_matrix( (conjMat[np.arange(i, conjMat.shape[0], args.osem), :]), dtype=np.float32) for i in range(args.osem) ] conjMats = [ chainer.utils.CooMatrix(p.data, p.row, p.col, p.shape) for p in conjMats ] # conjMat = sp.coo_matrix(conjMat, dtype = np.float32) # conjMat = chainer.utils.CooMatrix(conjMat.data, conjMat.row, conjMat.col, conjMat.shape) print("Conjugate matrix {} shape {}, thinned {} x {}".format( args.system_matrix, conjMat.shape, conjMats[0].shape, len(conjMats))) # setup updater print("Setting up data iterators...") planct_dataset = Dataset(path=args.planct_dir, baseA=args.HU_base, rangeA=args.HU_range, crop=(args.crop_height, args.crop_width), scale_to=args.scale_to, random=args.random_translate) planct_iter = chainer.iterators.SerialIterator(planct_dataset, args.batchsize, shuffle=True) mvct_dataset = Dataset(path=args.mvct_dir, baseA=args.HU_base, rangeA=args.HU_range, crop=(args.crop_height, args.crop_width), scale_to=args.scale_to, random=args.random_translate) mvct_iter = chainer.iterators.SerialIterator(mvct_dataset, args.batchsize, shuffle=True) data = prjData(args.sinogram, osem=args.osem) proj_iter = chainer.iterators.SerialIterator(data, args.batchsize, shuffle=False) # True updater = Updater(models=(seed, encoder, decoder, dis), iterator={ 'main': proj_iter, 'planct': planct_iter, 'mvct': mvct_iter }, optimizer={ 'main': optimizer_sd, 'enc': optimizer_enc, 'dec': optimizer_dec, 'dis': optimizer_dis }, device=args.gpu, params={ 'args': args, 'prMats': prMats, 'conjMats': conjMats }) # logging if args.epoch < 0: total_iter = -args.epoch * args.iter * math.ceil( len(data) / args.batchsize) else: total_iter = args.epoch * args.iter trainer = training.Trainer(updater, (total_iter, 'iteration'), out=args.out) log_interval = (50, 'iteration') log_keys_main = [] log_keys_dis = [] log_keys_grad = [ 'main/grad_sd', 'main/grad_gen', 'main/grad_sd_consistency', 'main/grad_gen_consistency', 'main/seed_diff' ] loss_main_list = [(args.lambda_sd, 'main/loss_sd'), (args.lambda_nn, 'main/loss_nn'), (args.lambda_ae1, 'main/loss_ae1'), (args.lambda_ae2, 'main/loss_ae2'), (args.lambda_tv, 'main/loss_tv'), (args.lambda_tvs, 'main/loss_tvs'), (args.lambda_reg, 'main/loss_reg'), (args.lambda_reg, 'main/loss_reg_ae')] for a, k in loss_main_list: if a > 0: log_keys_main.append(k) loss_dis_list = [(args.lambda_adv, 'main/loss_adv'), (args.lambda_advs, 'main/loss_advs'), (args.dis_freq, 'main/loss_dis'), (args.lambda_gan, 'main/loss_gan')] for a, k in loss_dis_list: if a > 0: log_keys_dis.append(k) log_keys = ['iteration'] + log_keys_main + log_keys_dis + log_keys_grad trainer.extend(extensions.observe_lr(), trigger=log_interval) trainer.extend(extensions.LogReport(keys=log_keys, trigger=log_interval)) trainer.extend(extensions.PrintReport(log_keys), trigger=log_interval) trainer.extend(extensions.ProgressBar(update_interval=10)) if extensions.PlotReport.available(): trainer.extend( extensions.PlotReport(log_keys_main, 'iteration', trigger=(100, 'iteration'), file_name='loss.png', postprocess=plot_log)) trainer.extend( extensions.PlotReport(log_keys_dis, 'iteration', trigger=(100, 'iteration'), file_name='loss_dis.png')) trainer.extend( extensions.PlotReport(log_keys_grad, 'iteration', trigger=(100, 'iteration'), file_name='loss_grad.png', postprocess=plot_log)) # trainer.extend(extensions.ParameterStatistics([seed,decoder])) ## very slow trainer.extend(CommandsExtension()) if args.snapinterval <= 0: args.snapinterval = total_iter if args.lambda_nn > 0: trainer.extend( extensions.dump_graph('main/loss_nn', out_name='gen.dot')) elif args.lambda_ae1 > 0: trainer.extend( extensions.dump_graph('main/loss_ae1', out_name='gen.dot')) # save models if args.use_enc: trainer.extend(extensions.snapshot_object( encoder, 'enc_{.updater.iteration}.npz'), trigger=(args.snapinterval, 'iteration')) trainer.extend(extensions.snapshot_object( optimizer_enc, 'opt_enc_{.updater.iteration}.npz'), trigger=(args.snapinterval, 'iteration')) if args.use_dis: trainer.extend(extensions.snapshot_object( dis, 'dis_{.updater.iteration}.npz'), trigger=(args.snapinterval, 'iteration')) trainer.extend(extensions.snapshot_object( optimizer_dis, 'opt_dis_{.updater.iteration}.npz'), trigger=(args.snapinterval, 'iteration')) # trainer.extend(extensions.dump_graph('main/loss_real', out_name='dis.dot')) trainer.extend(extensions.snapshot_object(decoder, 'dec_{.updater.iteration}.npz'), trigger=(args.snapinterval, 'iteration')) trainer.extend(extensions.snapshot_object( optimizer_dec, 'opt_dec_{.updater.iteration}.npz'), trigger=(args.snapinterval, 'iteration')) # save command line arguments os.makedirs(args.out, exist_ok=True) save_args(args, args.out) with open(os.path.join(args.out, "args.txt"), 'w') as fh: fh.write(" ".join(sys.argv)) trainer.run()
def main(): args = arguments() out = os.path.join(args.out, dt.now().strftime('%m%d_%H%M')) print(args) print("\nresults are saved under: ", out) save_args(args, out) 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() # cuda.cupy.cuda.set_allocator(cuda.cupy.cuda.MemoryPool().malloc) # 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, random=args.random_translate, forceSpacing=0) train_B_dataset = Dataset(path=os.path.join(args.root, 'trainB'), args=args, random=args.random_translate, forceSpacing=args.forceSpacing) test_A_dataset = Dataset(path=os.path.join(args.root, 'testA'), args=args, random=0, forceSpacing=0) test_B_dataset = Dataset(path=os.path.join(args.root, 'testB'), args=args, random=0, forceSpacing=args.forceSpacing) 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) if args.batch_size > 1: train_A_iter = chainer.iterators.MultiprocessIterator(train_A_dataset, args.batch_size, n_processes=3) train_B_iter = chainer.iterators.MultiprocessIterator(train_B_dataset, args.batch_size, n_processes=3) else: train_A_iter = chainer.iterators.SerialIterator( train_A_dataset, args.batch_size) train_B_iter = chainer.iterators.SerialIterator( train_B_dataset, args.batch_size) # setup models enc_x = net.Encoder(args) enc_y = enc_x if args.single_encoder else net.Encoder(args) dec_x = net.Decoder(args) dec_y = net.Decoder(args) dis_x = net.Discriminator(args) dis_y = net.Discriminator(args) dis_z = net.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) 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)) optimizer.setup(model) 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=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_d = [ 'opt_x/loss_real', 'opt_x/loss_fake', 'opt_y/loss_real', 'opt_y/loss_fake', 'opt_z/loss_x', 'opt_z/loss_y' ] 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.extend( ['opt_enc_x/loss_reg', 'opt_enc_y/loss_reg', '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: log_keys.extend(['opt_dec_x/loss_id', 'opt_dec_y/loss_id']) 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_gp', 'opt_y/loss_gp', 'opt_z/loss_gp']) 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.report_start < 1: if args.lambda_tv > 0: trainer.extend( extensions.dump_graph('opt_dec_y/loss_tv', out_name='dec.dot')) if args.lambda_reg > 0: trainer.extend( extensions.dump_graph('opt_enc_x/loss_reg', out_name='enc.dot')) 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')) 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')) ## visualisation vis_folder = os.path.join(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 }, device=args.gpu[0]), trigger=(args.vis_freq, 'iteration')) ## output filenames of training dataset with open(os.path.join(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(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(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 trainer.run()
def dump_chainer_info(self): """ returns chainer, CuPy and CuDNN version info """ chainer.print_runtime_info()
from chainer import print_runtime_info # Show Chainer's environment information print_runtime_info()
def main(): # command line argument parsing parser = argparse.ArgumentParser( description='Multi-Perceptron classifier/regressor') parser.add_argument('train', help='Path to csv file') parser.add_argument('--root', '-R', default="betti", help='Path to image files') parser.add_argument('--val', help='Path to validation csv file', required=True) parser.add_argument('--regress', '-r', action='store_true', help='set for regression, otherwise classification') parser.add_argument('--time_series', '-ts', action='store_true', help='set for time series data') parser.add_argument('--batchsize', '-b', type=int, default=10, help='Number of samples in each mini-batch') parser.add_argument('--layer', '-l', type=str, choices=['res5', 'pool5'], default='pool5', help='output layer of the pretrained ResNet') parser.add_argument('--fch', type=int, nargs="*", default=[], help='numbers of channels for the last fc layers') parser.add_argument('--cols', '-c', type=int, nargs="*", default=[1], help='column indices in csv of target variables') parser.add_argument('--epoch', '-e', type=int, default=300, help='Number of sweeps over the dataset to train') parser.add_argument('--snapshot', '-s', type=int, default=100, help='snapshot interval') parser.add_argument('--initmodel', '-i', help='Initialize the model from given file') parser.add_argument('--random', '-rt', type=int, default=1, help='random translation') parser.add_argument('--gpu', '-g', type=int, help='GPU ID (negative value indicates CPU)') parser.add_argument('--loaderjob', '-j', type=int, default=3, help='Number of parallel data loading processes') parser.add_argument('--outdir', '-o', default='result', help='Directory to output the result') parser.add_argument('--optimizer', '-op', choices=optim.keys(), default='Adam', help='optimizer') parser.add_argument('--resume', type=str, default=None, help='Resume the training from snapshot') parser.add_argument('--predict', '-p', action='store_true', help='prediction with a specified model') parser.add_argument('--tuning_rate', '-tr', type=float, default=0.1, help='learning rate for pretrained layers') parser.add_argument('--dropout', '-dr', type=float, default=0, help='dropout ratio for the FC layers') parser.add_argument('--cw', '-cw', type=int, default=128, help='crop image width') parser.add_argument('--ch', '-ch', type=int, default=128, help='crop image height') parser.add_argument('--weight_decay', '-w', type=float, default=1e-6, help='weight decay for regularization') parser.add_argument('--wd_norm', '-wn', choices=['none', 'l1', 'l2'], default='l2', help='norm of weight decay for regularization') parser.add_argument('--dtype', '-dt', choices=dtypes.keys(), default='fp32', help='floating point precision') args = parser.parse_args() args.outdir = os.path.join(args.outdir, dt.now().strftime('%m%d_%H%M')) # Enable autotuner of cuDNN chainer.config.autotune = True chainer.config.dtype = dtypes[args.dtype] chainer.print_runtime_info() # read csv file train = Dataset(args.root, args.train, cw=args.cw, ch=args.ch, random=args.random, regression=args.regress, time_series=args.time_series, cols=args.cols) test = Dataset(args.root, args.val, cw=args.cw, ch=args.ch, regression=args.regress, time_series=args.time_series, cols=args.cols) ## if not args.gpu: if chainer.cuda.available: args.gpu = 0 else: args.gpu = -1 print(args) save_args(args, args.outdir) if args.regress: accfun = F.mean_absolute_error lossfun = F.mean_squared_error args.chs = len(args.cols) else: accfun = F.accuracy lossfun = F.softmax_cross_entropy args.chs = max(train.chs, test.chs) if len(args.cols) > 1: print("\n\nClassification only works with a single target.\n\n") exit() # Set up a neural network to train model = L.Classifier(Resnet(args), lossfun=lossfun, accfun=accfun) # Set up an optimizer optimizer = optim[args.optimizer]() optimizer.setup(model) if args.weight_decay > 0: if args.wd_norm == 'l2': optimizer.add_hook(chainer.optimizer.WeightDecay( args.weight_decay)) elif args.wd_norm == 'l1': optimizer.add_hook(chainer.optimizer_hooks.Lasso( args.weight_decay)) # slow update for pretrained layers if args.optimizer in ['Adam']: for func_name in model.predictor.base._children: for param in model.predictor.base[func_name].params(): param.update_rule.hyperparam.alpha *= args.tuning_rate if args.initmodel: print('Load model from: ', args.initmodel) chainer.serializers.load_npz(args.initmodel, model) if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current model.to_gpu() # Copy the model to the GPU # select numpy or cupy xp = chainer.cuda.cupy if args.gpu >= 0 else np # train_iter = iterators.SerialIterator(train, args.batchsize, shuffle=True) # test_iter = iterators.SerialIterator(test, args.batchsize, repeat=False, shuffle=False) train_iter = iterators.MultithreadIterator(train, args.batchsize, shuffle=True, n_threads=args.loaderjob) test_iter = iterators.MultithreadIterator(test, args.batchsize, repeat=False, shuffle=False, n_threads=args.loaderjob) updater = training.StandardUpdater(train_iter, optimizer, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.outdir) frequency = args.epoch if args.snapshot == -1 else max(1, args.snapshot) log_interval = 1, 'epoch' val_interval = 20, 'epoch' # frequency/10, 'epoch' # trainer.extend(extensions.snapshot(filename='snapshot_epoch_{.updater.epoch}'), trigger=(frequency, 'epoch')) trainer.extend(extensions.snapshot_object(model, 'model_epoch_{.updater.epoch}'), trigger=(frequency, 'epoch')) trainer.extend(extensions.LogReport(trigger=log_interval)) trainer.extend(extensions.dump_graph('main/loss')) trainer.extend(extensions.Evaluator(test_iter, model, device=args.gpu), trigger=val_interval) if args.optimizer in ['Momentum', 'AdaGrad', 'RMSprop']: trainer.extend(extensions.observe_lr(), trigger=log_interval) trainer.extend(extensions.ExponentialShift('lr', 0.5), trigger=(args.epoch / 5, 'epoch')) elif args.optimizer in ['Adam']: trainer.extend(extensions.observe_lr(), trigger=log_interval) trainer.extend(extensions.ExponentialShift("alpha", 0.5, optimizer=optimizer), trigger=(args.epoch / 5, 'epoch')) if extensions.PlotReport.available(): trainer.extend( extensions.PlotReport(['main/loss', 'validation/main/loss'], 'epoch', file_name='loss.png')) trainer.extend( extensions.PlotReport( ['main/accuracy', 'validation/main/accuracy'], 'epoch', file_name='accuracy.png')) trainer.extend(extensions.PrintReport([ 'epoch', 'main/loss', 'main/accuracy', 'validation/main/loss', 'validation/main/accuracy', 'elapsed_time', 'lr' ]), trigger=log_interval) trainer.extend(extensions.ProgressBar(update_interval=10)) if args.resume: chainer.serializers.load_npz(args.resume, trainer) # ChainerUI #trainer.extend(CommandsExtension()) trainer.extend(extensions.LogReport(trigger=log_interval)) if not args.predict: trainer.run() ## prediction print("predicting: {} entries...".format(len(test))) test_iter = iterators.SerialIterator(test, args.batchsize, repeat=False, shuffle=False) converter = concat_examples idx = 0 with open(os.path.join(args.outdir, 'result.txt'), 'w') as output: for batch in test_iter: x, t = converter(batch, device=args.gpu) with chainer.using_config('train', False): with chainer.function.no_backprop_mode(): if args.regress: y = model.predictor(x).data if args.gpu > -1: y = xp.asnumpy(y) t = xp.asnumpy(t) y = y * test.std + test.mean t = t * test.std + test.mean else: y = F.softmax(model.predictor(x)).data if args.gpu > -1: y = xp.asnumpy(y) t = xp.asnumpy(t) for i in range(y.shape[0]): output.write(os.path.basename(test.ids[idx])) if (len(t.shape) > 1): for j in range(t.shape[1]): output.write(",{}".format(t[i, j])) output.write(",{}".format(y[i, j])) else: output.write(",{}".format(t[i])) output.write(",{}".format(np.argmax(y[i, :]))) for yy in y[i]: output.write(",{0:1.5f}".format(yy)) output.write("\n") idx += 1
def main(): # command line argument parsing parser = argparse.ArgumentParser(description='Digraph Embedding') parser.add_argument('input', help='Path to the digraph description file') parser.add_argument( '--validation', '-val', default=None, help='Path to the digraph description file for validation') parser.add_argument('--coordinates', '-c', help='Path to the coordinate file for initialization') parser.add_argument('--batchsize_edge', '-be', type=int, default=100, help='Number of samples in each edge mini-batch') parser.add_argument('--batchsize_anchor', '-ba', type=int, default=-1, help='Number of samples in each anchor mini-batch') parser.add_argument( '--batchsize_vert', '-bv', type=int, default=-1, help= 'Number of samples in each vertex mini-batch (used for sampling negative edges)' ) parser.add_argument( '--batchsize_negative', '-bn', type=int, default=0, help= 'Number of negative edges sampled for each vertex mini-batch (positive: exact negative edge sampling, negative: random sampling to approximate negative edges)' ) parser.add_argument('--vertex_offset', type=int, default=0, help='the smallest index of vertices') 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('--dim', '-d', type=int, default=2, help='Embedding dimension') parser.add_argument('--dag', type=float, default=0, help='0:non-acyclic, 1:acyclic') parser.add_argument('--margin', '-m', type=float, default=0.01, help='margin for the metric boundary') parser.add_argument('--weight_decay', '-wd', type=float, default=0, help='weight decay for regularization on coordinates') parser.add_argument('--wd_norm', '-wn', choices=['l1', 'l2'], default='l2', help='norm of weight decay for regularization') parser.add_argument('--learning_rate', '-lr', type=float, default=5e-2, help='learning rate') parser.add_argument('--learning_rate_drop', '-ld', type=int, default=5, help='how many times to half learning rate') # parser.add_argument('--lambda_super_neg', '-lsn', type=float, default=0, # help='Super negative samples') parser.add_argument('--lambda_pos', '-lp', type=float, default=1, help='weight for loss for positive edges') parser.add_argument('--lambda_neg', '-ln', type=float, default=1, help='weight for loss for negative edges') parser.add_argument( '--lambda_anchor', '-la', type=float, default=1, help= 'anchor should reside in the disk. if set to 0, anchors are fixed to the centre of the spheres' ) parser.add_argument('--lambda_uniform_radius', '-lur', type=float, default=0, help='all radiuses should be similar') parser.add_argument('--outdir', '-o', default='result', help='Directory to output the result') parser.add_argument('--optimizer', '-op', choices=optim.keys(), default='Adam', help='optimizer') parser.add_argument('--vis_freq', '-vf', type=int, default=-1, help='evaluation frequency in iteration') parser.add_argument('--mpi', action='store_true', help='parallelise with MPI') parser.add_argument('--reconstruct', '-r', action='store_true', help='reconstruct graph during evaluation') parser.add_argument('--plot', '-p', action='store_true', help='plot result (dim=2 only)') # parser.add_argument('--training', '-t', action='store_false',help='reconstruct graph') args = parser.parse_args() # default batchsize if args.batchsize_anchor < 0: args.batchsize_anchor = 10 * args.batchsize_edge if args.batchsize_vert < 0: if args.batchsize_negative == 0: args.batchsize_vert = 10 * args.batchsize_edge else: args.batchsize_vert = args.batchsize_edge args.outdir = os.path.join(args.outdir, dt.now().strftime('%m%d_%H%M')) save_args(args, args.outdir) chainer.config.autotune = True vert, pos_edge = read_graph(args.input, args.vertex_offset) vnum = np.max(vert) + 1 ## ChainerMN if args.mpi: import chainermn if args.gpu >= 0: comm = chainermn.create_communicator() chainer.cuda.get_device(comm.intra_rank).use() else: comm = chainermn.create_communicator('naive') if comm.rank == 0: primary = True print(args) chainer.print_runtime_info() print("#edges {}, #vertices {}".format(len(pos_edge), len(vert))) else: primary = False print("process {}".format(comm.rank)) else: primary = True print(args) chainer.print_runtime_info() print("#edges {}, #vertices {}".format(len(pos_edge), len(vert))) if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() # data edge_iter = iterators.SerialIterator(datasets.TupleDataset( pos_edge[:, 0], pos_edge[:, 1]), args.batchsize_edge, shuffle=True) vert_iter = iterators.SerialIterator(datasets.TupleDataset(vert), args.batchsize_vert, shuffle=True) anchor_iter = iterators.SerialIterator(datasets.TupleDataset(vert), args.batchsize_anchor, shuffle=True) graph = nx.from_edgelist(pos_edge, nx.DiGraph()) if args.validation and primary: val_vert, val_edge = read_graph(args.validation, args.vertex_offset) val_graph = nx.from_edgelist(val_edge, nx.DiGraph()) print("validation #edges {}, #vertices {}".format( len(val_edge), len(val_vert))) else: val_graph = graph if args.vis_freq < 0: args.vis_freq = int(len(pos_edge) * args.epoch / 10) # initial embedding if args.coordinates: coords = np.loadtxt(args.coordinates, delimiter=",") else: coords = np.zeros((vnum, 1 + 2 * args.dim)) # anchor = centre X = 2 * np.random.rand(vnum, args.dim) - 1 coords[:, 1:args.dim + 1] = X coords[:, args.dim + 1:] = X # the first coordinate corresponds to the radius r=0.1 coords[:, 0] = 0.1 coords = L.Parameter(coords) # set up an optimizer def make_optimizer(model): if args.optimizer in [ 'SGD', 'Momentum', 'CMomentum', 'AdaGrad', 'RMSprop', 'NesterovAG', 'LBFGS' ]: optimizer = optim[args.optimizer](lr=args.learning_rate) elif args.optimizer in ['AdaDelta']: optimizer = optim[args.optimizer]() elif args.optimizer in ['Adam', 'AdaBound', 'Eve']: optimizer = optim[args.optimizer]( alpha=args.learning_rate, weight_decay_rate=args.weight_decay) if args.mpi: optimizer = chainermn.create_multi_node_optimizer(optimizer, comm) optimizer.setup(model) return optimizer opt = make_optimizer(coords) if args.weight_decay > 0 and (not args.optimizer in ['Adam', 'AdaBound', 'Eve']): if args.wd_norm == 'l2': opt.add_hook(chainer.optimizer_hooks.WeightDecay( args.weight_decay)) else: opt.add_hook(chainer.optimizer_hooks.Lasso(args.weight_decay)) if args.gpu >= 0: coords.to_gpu() updater = Updater( models=coords, iterator={ 'main': edge_iter, 'vertex': vert_iter, 'anchor': anchor_iter }, optimizer={'main': opt}, device=args.gpu, # converter=convert.ConcatWithAsyncTransfer(), params={ 'args': args, 'graph': graph }) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.outdir) if primary: log_interval = 20, 'iteration' log_keys = [ 'iteration', 'lr', 'elapsed_time', 'main/loss_pos', 'main/loss_neg', 'main/loss_anc' ] if args.validation: log_keys.extend( ['myval/prc', 'myval/rec', 'myval/f1', 'myval/anc']) if args.lambda_uniform_radius > 0: log_keys.append('main/loss_rad') trainer.extend(extensions.observe_lr('main'), trigger=log_interval) trainer.extend( extensions.LogReport(keys=log_keys, trigger=log_interval)) # trainer.extend(extensions.LogReport(keys=log_keys, trigger=log_interval)) trainer.extend(extensions.PrintReport(log_keys), trigger=log_interval) # trainer.extend(extensions.PrintReport(log_keys), trigger=(1, 'iteration')) if extensions.PlotReport.available(): trainer.extend( extensions.PlotReport(log_keys[3:], 'epoch', file_name='loss.png', postprocess=plot_log)) trainer.extend(extensions.ProgressBar(update_interval=10)) trainer.extend(extensions.LogReport(trigger=log_interval)) trainer.extend(extensions.snapshot_object(opt, 'opt{.updater.epoch}.npz'), trigger=(args.epoch, 'epoch')) if args.vis_freq > 0: trainer.extend(Evaluator({'main': edge_iter}, coords, params={ 'args': args, 'graph': val_graph }, device=args.gpu), trigger=(args.vis_freq, 'iteration')) # trainer.extend(extensions.ParameterStatistics(coords)) # ChainerUI save_args(args, args.outdir) if args.optimizer in [ 'Momentum', 'CMomentum', 'AdaGrad', 'RMSprop', 'NesterovAG' ]: trainer.extend(extensions.ExponentialShift('lr', 0.5, optimizer=opt), trigger=(args.epoch / args.learning_rate_drop, 'epoch')) elif args.optimizer in ['Adam', 'AdaBound', 'Eve']: trainer.extend(extensions.ExponentialShift("alpha", 0.5, optimizer=opt), trigger=(args.epoch / args.learning_rate_drop, 'epoch')) # if args.training: trainer.run() # result if primary: # save DAG data file if (args.gpu > -1): dat = coords.xp.asnumpy(coords.W.data) else: dat = coords.W.data if args.lambda_anchor == 0: # anchor = centre dat[:, 1:(args.dim + 1)] = dat[:, (args.dim + 1):] redge = reconstruct(dat, dag=args.dag) np.savetxt(os.path.join(args.outdir, "original.csv"), pos_edge, fmt='%i', delimiter=",") np.savetxt(os.path.join(args.outdir, "reconstructed.csv"), redge, fmt='%i', delimiter=",") np.savetxt(os.path.join(args.outdir, "coords.csv"), dat, fmt='%1.5f', delimiter=",") f1, prc, rec, acc = compare_graph( val_graph, nx.from_edgelist(redge, nx.DiGraph())) if args.plot: plot_digraph(pos_edge, os.path.join(args.outdir, "original.png")) plot_digraph(redge, os.path.join(args.outdir, "reconstructed.png")) plot_disks(dat, os.path.join(args.outdir, "plot.png")) with open(os.path.join(args.outdir, "args.txt"), 'w') as fh: fh.write(" ".join(sys.argv)) fh.write( f"f1: {f1}, precision: {prc}, recall: {rec}, accuracy: {acc}")
def main(): parser = argparse.ArgumentParser(description='A Neural Algorithm of Artistic Style') parser.add_argument('input', default=None, help='Original image') parser.add_argument('--style', '-s', default=None, help='Style image') parser.add_argument('--reduced_image', '-r', default=None, help='reduced contents image') parser.add_argument('--init_image', '-i', default=None, help="start optimisation using this image, otherwise start with a random image") parser.add_argument('--out', '-o', default='result', help='Output directory') parser.add_argument('--gpu', '-g', default=0, type=int, help='GPU ID (negative value indicates CPU)') parser.add_argument('--iter', default=5000, type=int, help='number of iterations') parser.add_argument('--vis_freq', '-vf', default=500, type=int, help='image output interval') parser.add_argument('--ksize', '-k', default=4, type=int, help='kernel size for reduction') parser.add_argument('--lr', default=4.0, type=float, help='learning rate') parser.add_argument('--lambda_tv', '-ltv', default=1, type=float, help='weight of total variation regularization') parser.add_argument('--lambda_feat', '-lf', default=0.1, type=float, help='weight for the original shape; increase to retain it') parser.add_argument('--lambda_rfeat', '-lrf', default=0.1, type=float, help='weight for the reduce shape; increase to retain it') parser.add_argument('--lambda_style', '-ls', default=1.0, type=float, help='weight for the style') parser.add_argument('--median_filter', '-f', default=3, type=int, help='kernel size of the median filter applied to the output') parser.add_argument('--removebg', '-nbg', action='store_true', help="remove background specified by alpha in png") parser.add_argument('--crop_width', '-cw', type=int, default=0) parser.add_argument('--crop_height', '-ch', type=int, default=0) ## dicom related parser.add_argument('--image_size', default=448, type=int) parser.add_argument('--CT_base', '-ctb', type=int, default=-250) parser.add_argument('--CT_range', '-ctr', type=int, default=500) parser.add_argument('--CT_A_scale', type=float, default=1.0) parser.add_argument('--CT_B_scale', type=float, default=2.148/0.7634) args = parser.parse_args() chainer.config.autotune = True chainer.print_runtime_info() if not args.style: args.lambda_style = 0 if not args.reduced_image: args.lambda_rfeat = 0 print(args) os.makedirs(args.out, exist_ok=True) if args.gpu >= 0: cuda.get_device(args.gpu).use() xp = cuda.cupy else: print("runs desperately slowly without a GPU!") xp = np ## use pretrained VGG for feature extraction nn = L.VGG16Layers() ## load images if os.path.splitext(args.input)[1] == '.dcm': image = load_dicom(args.input, args.CT_base, args.CT_range, args.image_size, args.CT_A_scale) image = xp.tile(image,(1,3,1,1)) style = load_dicom(args.style, args.CT_base, args.CT_range, args.image_size, args.CT_B_scale) style = xp.tile(style,(1,3,1,1)) else: image = xp.asarray(load_image(args.input,size=(args.crop_width,args.crop_height),removebg=args.removebg)) if args.lambda_style > 0: style = xp.asarray(load_image(args.style,size=(image.shape[3],image.shape[2]), removebg=args.removebg)) bg = Image.open(args.input).convert('RGBA') if args.crop_height>0 and args.crop_width>0: bg = bg.resize((image.shape[3],image.shape[2]),Image.LANCZOS) bg = np.array(bg) bg[:,:,3] = 255-bg[:,:,3] bkgnd = Image.fromarray(bg) ## initial image: original or random if args.init_image: img_gen = load_image(args.init_image,size=(image.shape[3],image.shape[2]), removebg=args.removebg) else: img_gen = xp.random.uniform(-20,20,image.shape).astype(np.float32) print("input image:", image.shape, xp.min(image), xp.max(image)) ## image to be generated img_gen = L.Parameter(img_gen) if args.gpu>=0: nn.to_gpu() img_gen.to_gpu() optimizer = optimizers.Adam(alpha=args.lr) optimizer.setup(img_gen) ## compute style matrix layers = ["conv1_2","conv2_2","conv3_3","conv4_3"] with chainer.using_config('train', False): feature = nn(Variable(image),layers=layers) if args.lambda_style > 0: feature_s = nn(Variable(style),layers=layers) gram_s = { k:gram_matrix(feature_s[k]) for k in layers} else: gram_s = None if args.lambda_rfeat > 0: reduced_image = xp.asarray(load_image(args.reduced_image,size=(image.shape[3],image.shape[2]),removebg=args.removebg)) reduced_feature = nn(F.average_pooling_2d(Variable(reduced_image),args.ksize),layers=layers) else: reduced_feature = None # modify loss weights according to the feature vector size args.lambda_rfeat /= args.ksize ** 2 # setup updater dummy_iterator = chainer.iterators.SerialIterator(np.array([1]),1) updater = Updater( models=(img_gen,nn), iterator=dummy_iterator, optimizer=optimizer, # converter=convert.ConcatWithAsyncTransfer(), device=args.gpu, params={'args': args, 'image': image, 'reduced_feature': reduced_feature, 'bkgnd': bkgnd, 'feature': feature, 'gram_s': gram_s, 'layers': layers} ) trainer = training.Trainer(updater, (args.iter, 'iteration'), out=args.out) log_interval = (100, 'iteration') log_keys = ['iteration','lr','main/loss_tv','main/loss_f','main/loss_s','main/loss_r'] trainer.extend(extensions.observe_lr(), trigger=log_interval) trainer.extend(extensions.LogReport(keys=log_keys, trigger=log_interval)) trainer.extend(extensions.PrintReport(log_keys), trigger=log_interval) trainer.extend(extensions.ProgressBar(update_interval=10)) if extensions.PlotReport.available(): trainer.extend(extensions.PlotReport( log_keys[2:], 'iteration', trigger=(1000, 'iteration'), file_name='loss.png')) trainer.run()