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()
