def build_model(self):
# Define generators and discriminators
if self.whichG=='normal':
self.G = net.Generator_makeup(self.g_conv_dim, self.g_repeat_num)
if self.whichG=='branch':
self.G = net.Generator_branch(self.g_conv_dim, self.g_repeat_num)
for i in self.cls:
setattr(self, "D_" + i, net.Discriminator(self.img_size, self.d_conv_dim, self.d_repeat_num, self.norm))
self.criterionL1 = torch.nn.L1Loss()
self.criterionL2 = torch.nn.MSELoss()
self.criterionGAN = GANLoss(use_lsgan=True, tensor =torch.cuda.FloatTensor)
self.vgg=models.vgg16(pretrained=True)
# Optimizers
self.g_optimizer = torch.optim.Adam(self.G.parameters(), self.g_lr, [self.beta1, self.beta2])
for i in self.cls:
setattr(self, "d_" + i + "_optimizer", \
torch.optim.Adam(filter(lambda p: p.requires_grad, getattr(self, "D_" + i).parameters()), \
self.d_lr, [self.beta1, self.beta2]))
# Weights initialization
self.G.apply(self.weights_init_xavier)
for i in self.cls:
getattr(self, "D_" + i).apply(self.weights_init_xavier)
if torch.cuda.is_available():
self.G.cuda()
self.vgg.cuda()
for i in self.cls:
getattr(self, "D_" + i).cuda()
def build_model(self):
# Define generators and discriminators
if self.whichG == 'normal':
self.G = net.Generator_makeup(self.g_conv_dim, self.g_repeat_num)
if self.whichG == 'branch':
self.G = net.Generator_branch(self.g_conv_dim, self.g_repeat_num)
for i in self.cls:
setattr(
self, "D_" + i,
net.Discriminator(self.img_size, self.d_conv_dim,
self.d_repeat_num, self.norm))
self.criterionL1 = torch.nn.L1Loss()
self.criterionL2 = torch.nn.MSELoss()
self.criterionGAN = GANLoss(use_lsgan=True,
tensor=torch.cuda.FloatTensor)
self.vgg = net.VGG()
self.vgg.load_state_dict(torch.load('addings/vgg_conv.pth'))
# self.vgg = models.vgg19_bn(pretrained=True)
# Optimizers
self.g_optimizer = torch.optim.Adam(self.G.parameters(), self.g_lr,
[self.beta1, self.beta2])
for i in self.cls:
setattr(self, "d_" + i + "_optimizer", \
torch.optim.Adam(filter(lambda p: p.requires_grad, getattr(self, "D_" + i).parameters()), \
self.d_lr, [self.beta1, self.beta2]))
# Weights initialization
self.G.apply(self.weights_init_xavier)
for i in self.cls:
getattr(self, "D_" + i).apply(self.weights_init_xavier)
# Print networks
self.print_network(self.G, 'G')
for i in self.cls:
self.print_network(getattr(self, "D_" + i), "D_" + i)
"""
if torch.cuda.device_count() > 1:
self.G = torch.nn.DataParallel(self.G)
self.vgg = torch.nn.DataParallel(self.vgg)
for i in self.cls:
setattr(self, "D_" + i, torch.nn.DataParallel(getattr(self, "D_" + i)))
self.G.to(self.device)
self.vgg.to(self.device)
for i in self.cls:
getattr(self, "D_" + i).to(self.device)
"""
if torch.cuda.is_available():
self.G.cuda()
self.vgg.cuda()
for i in self.cls:
getattr(self, "D_" + i).cuda()
def build_model(self):
# Define generators and discriminators
self.G_A = net.Generator(self.g_conv_dim, self.g_repeat_num)
self.G_B = net.Generator(self.g_conv_dim, self.g_repeat_num)
self.D_A = net.Discriminator(self.img_size, self.d_conv_dim,
self.d_repeat_num)
self.D_B = net.Discriminator(self.img_size, self.d_conv_dim,
self.d_repeat_num)
self.criterionL1 = torch.nn.L1Loss()
self.criterionGAN = GANLoss(use_lsgan=True,
tensor=torch.cuda.FloatTensor)
# Optimizers
self.g_optimizer = torch.optim.Adam(
itertools.chain(self.G_A.parameters(), self.G_B.parameters()),
self.g_lr, [self.beta1, self.beta2])
self.d_A_optimizer = torch.optim.Adam(
filter(lambda p: p.requires_grad, self.D_A.parameters()),
self.d_lr, [self.beta1, self.beta2])
self.d_B_optimizer = torch.optim.Adam(
filter(lambda p: p.requires_grad, self.D_B.parameters()),
self.d_lr, [self.beta1, self.beta2])
self.G_A.apply(self.weights_init_xavier)
self.D_A.apply(self.weights_init_xavier)
self.G_B.apply(self.weights_init_xavier)
self.D_B.apply(self.weights_init_xavier)
# Print networks
# self.print_network(self.E, 'E')
self.print_network(self.G_A, 'G_A')
self.print_network(self.D_A, 'D_A')
self.print_network(self.G_B, 'G_B')
self.print_network(self.D_B, 'D_B')
if torch.cuda.is_available():
self.G_A.cuda()
self.G_B.cuda()
self.D_A.cuda()
self.D_B.cuda()
def build_model(self):
# Define generators and discriminators
self.E = network.Encoder(self.e_conv_dim)
self.G = network.Generator(self.g_conv_dim)
for i in self.cls:
setattr(
self, "D_" + i,
net.Discriminator(self.img_size, self.d_conv_dim,
self.d_repeat_num, self.norm))
# Define vgg for perceptual loss
self.vgg = net.VGG()
self.vgg.load_state_dict(torch.load('addings/vgg_conv.pth'))
# Define loss
self.criterionL1 = torch.nn.L1Loss()
self.criterionL2 = torch.nn.MSELoss()
self.criterionGAN = GANLoss(use_lsgan=True,
tensor=torch.cuda.FloatTensor)
# Optimizers
self.e_optimizer = torch.optim.Adam(self.E.parameters(), self.e_lr,
[self.beta1, self.beta2])
self.g_optimizer = torch.optim.Adam(self.G.parameters(), self.g_lr,
[self.beta1, self.beta2])
for i in self.cls:
setattr(self, "d_" + i + "_optimizer", \
torch.optim.Adam(filter(lambda p: p.requires_grad, getattr(self, "D_" + i).parameters()), \
self.d_lr, [self.beta1, self.beta2]))
# Weights initialization
self.E.apply(self.weights_init_xavier)
self.G.apply(self.weights_init_xavier)
for i in self.cls:
getattr(self, "D_" + i).apply(self.weights_init_xavier)
# Print networks
self.print_network(self.E, 'E')
self.print_network(self.G, 'G')
for i in self.cls:
self.print_network(getattr(self, "D_" + i), "D_" + i)
if torch.cuda.is_available():
self.E.cuda()
self.G.cuda()
self.vgg.cuda()
for i in self.cls:
getattr(self, "D_" + i).cuda()
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()
import numpy as np
from chainer import Variable
import chainer.computational_graph as c
import net
gen = net.Generator(784, 20, 500)
dis = net.Discriminator(784, 500)
x_real = np.empty((1, 784), dtype=np.float32)
z = Variable(np.asarray(gen.make_hidden(1)))
y_real = dis(x_real)
x_fake = gen(z)
y_fake = dis(x_fake)
g = c.build_computational_graph([y_real, x_fake, y_fake])
with open('graph.dot', 'w') as o:
o.write(g.dump())
help='dataset file path')
parser.add_argument('--size',
'-s',
default=96,
type=int,
choices=[48, 96],
help='image size')
args = parser.parse_args()
image_size = args.size
if image_size == 48:
gen_model = net.Generator48()
dis_model = net.Discriminator48()
else:
gen_model = net.Generator()
dis_model = net.Discriminator()
optimizer_gen = optimizers.Adam(alpha=0.0002, beta1=0.5)
optimizer_gen.setup(gen_model)
optimizer_gen.add_hook(chainer.optimizer.WeightDecay(0.00001))
optimizer_dis = optimizers.Adam(alpha=0.0002, beta1=0.5)
optimizer_dis.setup(dis_model)
optimizer_dis.add_hook(chainer.optimizer.WeightDecay(0.00001))
if args.input != None:
serializers.load_hdf5(args.input + '.gen.model', gen_model)
serializers.load_hdf5(args.input + '.gen.state', optimizer_gen)
serializers.load_hdf5(args.input + '.dis.model', dis_model)
serializers.load_hdf5(args.input + '.dis.state', optimizer_dis)
if args.out_image_dir != None:
serializers.load_npz(modelfn, dec)
if args.gpu >= 0:
enc.to_gpu()
dec.to_gpu()
xp = enc.xp
is_AE = True
else:
print("Specify a learnt model.")
exit()
## prepare networks for analysis
if args.output_analysis:
if is_AE:
enc_y = net.Encoder(args)
dec_x = net.Decoder(args)
dis_x = net.Discriminator(args)
dis_y = net.Discriminator(args)
models = {
'enc_y': enc_y,
'dec_x': dec_x,
'dis_x': dis_x,
'dis_y': dis_y
}
else:
gen_f = net.Generator(args)
dis_y = net.Discriminator(args)
dis_x = net.Discriminator(args)
models = {'gen_f': gen_f, 'dis_x': dis_x, 'dis_y': dis_y}
for e in models:
path = args.load_models.replace('gen_g', e)
path = path.replace('enc_x', e)
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 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 = easydict.EasyDict({
# "dataroot": "/mnt/gold/users/s18150/mywork/pytorch/data/gan",
"dataroot": "/mnt/gold/users/s18150/mywork/pytorch/data",
"save_dir": "./",
"prefix": "test",
"workers": 8,
"batch_size": 128,
"image_size": 32,
# "image_size": 28,
# "nc": 3,
"nc": 1,
"nz": 100,
"ngf": 32,
"ndf": 32,
# "ngf": 28,
# "ndf": 64,
"epochs": 1,
"lr": 0.0002,
"beta1": 0.5,
"gpu": 7,
"use_cuda": True,
"feature_matching": True,
"mini_batch": True,
"iters": 50000,
"label_batch_size": 100,
"unlabel_batch_size": 100,
"test_batch_size": 10,
"out_dir": './result',
"log_interval": 500,
"label_num": 20
})
manualSeed = 999
np.random.seed(manualSeed)
random.seed(manualSeed)
torch.manual_seed(manualSeed)
device = torch.device(
'cuda:{}'.format(args.gpu) if args.use_cuda else 'cpu')
# transform = transforms.Compose([
# transforms.Resize(args.image_size),
# transforms.CenterCrop(args.image_size),
# transforms.ToTensor(),
# transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
# ])
#
# dataset = dset.ImageFolder(root=args.dataroot, transform=transform)
# dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size,
# shuffle=True, num_workers=args.workers)
data_iterators = dataset.get_iters(root_path=args.dataroot,
l_batch_size=args.label_batch_size,
ul_batch_size=args.unlabel_batch_size,
test_batch_size=args.test_batch_size,
workers=args.workers,
n_labeled=args.label_num)
trainloader_label = data_iterators['labeled']
trainloader_unlabel = data_iterators['unlabeled']
testloader = data_iterators['test']
# Generator用のモデルのインスタンス作成
netG = net.Generator(args.nz, args.ngf, args.nc).to(device)
# Generator用のモデルの初期値を設定
netG.apply(net.weights_init)
# Discriminator用のモデルのインスタンス作成
netD = net.Discriminator(args.nc, args.ndf, device, args.batch_size,
args.mini_batch).to(device)
# Discriminator用のモデルの初期値を設定
netD.apply(net.weights_init)
# BCE Loss classのインスタンスを作成
criterionD = nn.CrossEntropyLoss()
# criterionD = nn.BCELoss()
if args.feature_matching is True:
criterionG = nn.MSELoss(reduction='elementwise_mean')
else:
criterionG = nn.BCELoss()
# Generatorに入力するノイズをバッチごとに作成 (バッチ数は64)
# これはGeneratorの結果を描画するために使用する
fixed_noise = torch.randn(64, args.nz, 1, 1, device=device)
# 最適化関数のインスタンスを作成
optimizerD = optim.Adam(netD.parameters(),
lr=args.lr,
betas=(args.beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(),
lr=args.lr,
betas=(args.beta1, 0.999))
logger = TrainLogger(args)
r = run.NNRun(netD, netG, optimizerD, optimizerG, criterionD, criterionG,
device, fixed_noise, logger, args)
# 学習
# r.train(dataloader)
r.train(trainloader_label, trainloader_unlabel, testloader)
def gan_training(args, train, test):
# These iterators load the images with subprocesses running in parallel to
# the training/validation.
if args.loaderjob:
train_iter = chainer.iterators.MultiprocessIterator(
train, args.batchsize, n_processes=args.loaderjob)
else:
train_iter = chainer.iterators.SerialIterator(train, args.batchsize)
# Prepare Texture GAN model, defined in net.py
gen = net.Generator(args.dimz)
dis = net.Discriminator()
if args.gpu >= 0:
cuda.get_device(args.gpu).use()
gen.to_gpu()
dis.to_gpu()
xp = np if args.gpu < 0 else cuda.cupy
opt_gen = make_optimizer(gen, args, alpha=args.alpha)
opt_dis = make_optimizer(dis, args, alpha=args.alpha)
# Updater
updater = GAN_Updater(models=(gen, dis),
iterator=train_iter,
optimizer={
'gen': opt_gen,
'dis': opt_dis
},
device=args.gpu)
trainer = training.Trainer(updater, (args.iteration, 'iteration'),
out=args.out)
snapshot_interval = (args.snapshot_interval), 'iteration'
visualize_interval = (args.visualize_interval), 'iteration'
log_interval = (args.log_interval), 'iteration'
# Be careful to pass the interval directly to LogReport
# (it determines when to emit log rather than when to read observations)
trainer.extend(extensions.LogReport(trigger=log_interval))
trainer.extend(
extensions.PlotReport(['gen/loss', 'dis/loss'],
trigger=log_interval,
file_name='plot.png'))
trainer.extend(extensions.PrintReport(
['epoch', 'iteration', 'gen/loss', 'dis/loss']),
trigger=log_interval)
trainer.extend(extensions.ProgressBar(update_interval=10))
trainer.extend(extensions.snapshot(), trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
gen, 'gen_iteration_{.updater.iteration}.npz'),
trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
dis, 'dis_iteration_{.updater.iteration}.npz'),
trigger=snapshot_interval)
np.random.seed(0)
train_indices = np.random.randint(0, len(train),
args.rows * args.cols).tolist()
test_indices = np.random.randint(0, len(test),
args.rows * args.cols).tolist()
np.random.seed()
train_indices[-2] = len(train) - 3
train_indices[-3] = len(train) - 1
trainer.extend(visualizer.extension(train,
gen,
train_indices,
args,
'train',
rows=args.rows,
cols=args.cols),
trigger=visualize_interval)
trainer.extend(visualizer.extension(test,
gen,
test_indices,
args,
'test',
rows=args.rows,
cols=args.cols),
trigger=visualize_interval)
if args.adam_decay_iteration:
trainer.extend(extensions.ExponentialShift("alpha",
0.5,
optimizer=opt_gen),
trigger=(args.adam_decay_iteration, 'iteration'))
trainer.extend(extensions.ExponentialShift("alpha",
0.5,
optimizer=opt_dis),
trigger=(args.adam_decay_iteration, 'iteration'))
if args.resume:
chainer.serializers.load_npz(args.resume, trainer)
# Run the training
trainer.run()
def gan_training(args, train, test, model_path):
# These iterators load the images with subprocesses running in parallel to
# the training/validation.
if args.loaderjob:
train_iter = chainer.iterators.MultiprocessIterator(
train, args.batchsize, n_processes=args.loaderjob)
else:
train_iter = chainer.iterators.SerialIterator(train, args.batchsize)
# Prepare Flow and Texture GAN model, defined in net.py
gen_flow = net.FlowGenerator()
dis_flow = net.FlowDiscriminator()
gen_tex = net.Generator(args.dimz)
dis_tex = net.Discriminator()
serializers.load_npz(model_path["gen_flow"], gen_flow)
serializers.load_npz(model_path["dis_flow"], dis_flow)
serializers.load_npz(model_path["gen_tex"], gen_tex)
serializers.load_npz(model_path["dis_tex"], dis_tex)
if args.gpu >= 0:
cuda.get_device(args.gpu).use()
gen_flow.to_gpu()
dis_flow.to_gpu()
gen_tex.to_gpu()
dis_tex.to_gpu()
xp = np if args.gpu < 0 else cuda.cupy
opt_flow_gen = make_optimizer(gen_flow, args, alpha=1e-7)
opt_flow_dis = make_optimizer(dis_flow, args, alpha=1e-7)
opt_tex_gen = make_optimizer(gen_tex, args, alpha=1e-6)
opt_tex_dis = make_optimizer(dis_tex, args, alpha=1e-6)
# Updater
updater = GAN_Updater(models=(gen_flow, dis_flow, gen_tex, dis_tex),
iterator=train_iter,
optimizer={
'gen_flow': opt_flow_gen,
'dis_flow': opt_flow_dis,
'gen_tex': opt_tex_gen,
'dis_tex': opt_tex_dis
},
device=args.gpu,
C=args.C)
trainer = training.Trainer(updater, (args.iteration, 'iteration'),
out=args.out)
snapshot_interval = (args.snapshot_interval), 'iteration'
visualize_interval = (args.visualize_interval), 'iteration'
log_interval = (args.log_interval), 'iteration'
# Be careful to pass the interval directly to LogReport
# (it determines when to emit log rather than when to read observations)
trainer.extend(extensions.LogReport(trigger=log_interval))
trainer.extend(
extensions.PlotReport(
['gen_flow/loss', 'dis_flow/loss', 'gen_tex/loss', 'dis_tex/loss'],
trigger=log_interval,
file_name='plot.png'))
trainer.extend(extensions.PrintReport([
'epoch', 'iteration', 'gen_flow/loss', 'dis_flow/loss', 'gen_tex/loss',
'dis_tex/loss'
]),
trigger=log_interval)
trainer.extend(extensions.ProgressBar(update_interval=10))
trainer.extend(extensions.snapshot(), trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
gen_flow, 'gen_flow_iteration_{.updater.iteration}.npz'),
trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
dis_flow, 'dis_flow_iteration_{.updater.iteration}.npz'),
trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
gen_tex, 'gen_tex_iteration_{.updater.iteration}.npz'),
trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
dis_tex, 'dis_tex_iteration_{.updater.iteration}.npz'),
trigger=snapshot_interval)
trainer.extend(visualizer.extension(test, (gen_flow, gen_tex),
args,
rows=args.rows,
cols=args.cols),
trigger=visualize_interval)
if args.adam_decay_iteration:
trainer.extend(extensions.ExponentialShift("alpha",
0.5,
optimizer=opt_flow_gen),
trigger=(args.adam_decay_iteration, 'iteration'))
trainer.extend(extensions.ExponentialShift("alpha",
0.5,
optimizer=opt_flow_dis),
trigger=(args.adam_decay_iteration, 'iteration'))
trainer.extend(extensions.ExponentialShift("alpha",
0.5,
optimizer=opt_tex_gen),
trigger=(args.adam_decay_iteration, 'iteration'))
trainer.extend(extensions.ExponentialShift("alpha",
0.5,
optimizer=opt_tex_dis),
trigger=(args.adam_decay_iteration, 'iteration'))
if args.resume:
chainer.serializers.load_npz(args.resume, trainer)
# Run the training
trainer.run()
args = parser.parse_args()
device = torch.device('cuda')
if not os.path.exists(args.save_dir):
os.mkdir(args.save_dir)
if not os.path.exists(args.log_dir):
os.mkdir(args.log_dir)
writer = SummaryWriter(log_dir=args.log_dir)
decoder = net.decoder
vgg = net.vgg
discriminator = net.Discriminator()
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=args.lr)
discriminator.cuda()
vgg.load_state_dict(torch.load(args.vgg))
vgg = nn.Sequential(*list(vgg.children())[:31])
network = net.Net(vgg, decoder)
network.train()
network.to(device)
content_tf = train_transform()
style_tf = train_transform()
content_dataset = FlatFolderDataset(args.content_dir, content_tf)
style_dataset = FlatFolderDataset(args.style_dir, style_tf)
def main():
parser = argparse.ArgumentParser(description='WGAN(GP) MNIST')
parser.add_argument('--mode', '-m', type=str, default='WGAN',
help='WGAN or WGANGP')
parser.add_argument('--gpu', '-g', type=int, default=-1,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--epoch', '-e', type=int, default=100,
help='number of epochs to learn')
parser.add_argument('--batchsize', '-b', type=int, default=100,
help='learning minibatch size')
parser.add_argument('--optimizer', type=str, default='Adam',
help='optimizer')
parser.add_argument('--out', '-o', type=str, default='model',
help='path to the output directory')
parser.add_argument('--dimz', '-z', type=int, default=20,
help='dimention of encoded vector')
parser.add_argument('--n_dis', type=int, default=5,
help='dimention of encoded vector')
parser.add_argument('--seed', type=int, default=0,
help='Random seed of z at visualization stage')
parser.add_argument('--snapepoch', '-s', type=int, default=10,
help='number of epochs to snapshot')
parser.add_argument('--load_gen_model', type=str, default='',
help='load generator model')
parser.add_argument('--load_dis_model', type=str, default='',
help='load generator model')
args = parser.parse_args()
if not os.path.exists(args.out):
os.makedirs(args.out)
print(args)
gen = net.Generator(784, args.dimz, 500)
dis = net.Discriminator(784, 500)
if args.load_gen_model != '':
chainer.serializers.load_npz(args.load_gen_model, gen)
if args.load_dis_model != '':
chainer.serializers.load_npz(args.load_dis_model, dis)
if args.gpu >= 0:
cuda.get_device(args.gpu).use()
gen.to_gpu()
dis.to_gpu()
print('GPU {}'.format(args.gpu))
xp = np if args.gpu < 0 else cuda.cupy
if args.optimizer == 'Adam':
opt_gen = chainer.optimizers.Adam(alpha=0.0001, beta1=0.5, beta2=0.9)
opt_dis = chainer.optimizers.Adam(alpha=0.0001, beta1=0.5, beta2=0.9)
opt_gen.setup(gen)
opt_dis.setup(dis)
opt_dis.add_hook(WeightClipping(0.01))
elif args.optimizer == 'RMSprop':
opt_gen = chainer.optimizers.RMSprop(5e-5)
opt_dis = chainer.optimizers.RMSprop(5e-5)
opt_gen.setup(gen)
opt_gen.add_hook(chainer.optimizer.GradientClipping(1))
opt_dis.setup(dis)
opt_dis.add_hook(chainer.optimizer.GradientClipping(1))
opt_dis.add_hook(WeightClipping(0.01))
train, _ = chainer.datasets.get_mnist(withlabel=False)
train_iter = chainer.iterators.SerialIterator(train, args.batchsize, shuffle=True)
if args.mode == 'WGAN':
updater = WGANUpdater(
models=(gen, dis),
iterators={
'main': train_iter
},
optimizers={
'gen': opt_gen,
'dis': opt_dis
},
device=args.gpu,
params={
'batchsize': args.batchsize,
'n_dis': args.n_dis
})
elif args.mode == 'WGANGP':
updater = WGANGPUpdater(
models=(gen, dis),
iterators={
'main': train_iter
},
optimizers={
'gen': opt_gen,
'dis': opt_dis
},
device=args.gpu,
params={
'batchsize': args.batchsize,
'n_dis': args.n_dis,
'lam': 10
})
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
trainer.extend(extensions.dump_graph('was. dist'))
snapshot_interval = (args.snapepoch, 'epoch')
trainer.extend(extensions.snapshot(filename='snapshot_epoch_{.updater.epoch}.npz'),
trigger=snapshot_interval)
trainer.extend(extensions.PlotReport(['loss/gen'], 'epoch', file_name='generator.png'))
if args.mode == 'WGAN':
log_keys = ['epoch', 'was. dist', 'gen/loss', 'dis/loss']
elif args.mode == 'WGANGP':
log_keys = ['epoch', 'was. dist', 'grad. pen', 'gen/loss']
trainer.extend(extensions.LogReport(keys=log_keys))
trainer.extend(extensions.PrintReport(log_keys))
trainer.extend(extensions.ProgressBar(update_interval=10))
trainer.extend(out_generated_image(gen, 20, 20, args.seed, args.out),
trigger=(1, 'epoch'))
trainer.run()
train_dis = True
image_save_interval = 200000
model_save_interval = image_save_interval
out_image_row_num = 7
out_image_col_num = 14
if train_LSTM:
BATCH_SIZE /= seq_length
normer = args.size * args.size * 3 * 60
image_path = ['../../../trajectories/al5d']
np.random.seed(1241)
image_size = args.size
enc_model = [net.Encoder(density=8, size=image_size, latent_size=latent_size)]
gen_model = [
net.Generator(density=8, size=image_size, latent_size=latent_size)
]
dis_model = [net.Discriminator(density=8, size=image_size)]
for i in range(num_gpus - 1):
enc_model.append(copy.deepcopy(enc_model[0]))
gen_model.append(copy.deepcopy(gen_model[0]))
dis_model.append(copy.deepcopy(dis_model[0]))
enc_dis_model = net.Encoder(density=8,
size=image_size,
latent_size=latent_size)
gen_dis_model = net.Generator(density=8,
size=image_size,
latent_size=latent_size)
rnn_model = rnn_net.MDN_RNN(IN_DIM=latent_size + 5,
HIDDEN_DIM=300,
OUT_DIM=6,
NUM_MIXTURE=40)
def main():
args = parse_arg()
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
writer = SummaryWriter(os.path.join(args.save_dir, 'tb'))
if args.use_cuda:
device = torch.device('cuda')
else:
device = torch.device('cpu')
model = net.InPainting(args.use_cuda, mask_clipping=False).to(device)
model_dis = net.Discriminator().to(device)
model.train()
model_dis.train()
train_params = [
p for (n, p) in model.named_parameters() if 'fine_painter' not in n
]
optimizer = torch.optim.Adam(train_params, lr=args.lr)
optimizer_dis = torch.optim.Adam(model_dis.parameters(), lr=args.lr)
if args.resume:
print('Resume training from ', args.resume)
ckpt = torch.load(args.resume)
try:
model.load_state_dict(ckpt['ckpt'])
model.load_state_dict(ckpt['ckpt_dis'])
optimizer.load_state_dict(ckpt['optim'])
optimizer_dis.load_state_dict(ckpt['optim_dis'])
except Exception as e:
print(traceback.format_exc())
print('Missing keys')
model.load_state_dict(
{k: v
for k, v in ckpt['ckpt'].items() if 'encoder' in k},
strict=False)
if args.fix_mask:
print('fix mask prediction')
for p in model.encoder.parameters():
p.requires_grad = False
elif args.fix_recon:
print('fix painter')
for p in model.painter.parameters():
p.requires_grad = False
loss = torch.nn.L1Loss()
loss_bce = torch.nn.BCELoss()
ssim_window = ssim.create_window(11, 3).to(device)
#data = dataset.Dataset(args.path)
data = dataset.LargeScaleWatermarkDataset(
folder_origin=os.path.join(args.path, args.path_origin),
folder_watermarked=os.path.join(args.path, args.path_wm),
anno_file=os.path.join(args.path, args.path_anno))
train_loader = torch.utils.data.DataLoader(dataset=data,
batch_size=args.batchsize,
shuffle=True,
num_workers=4)
try:
batch_per_epoch = len(train_loader)
best_loss = 100
for i in range(args.epochs):
epoch_loss = 0
print('Epoch %d' % i)
for j, item in enumerate(train_loader):
img_raw, img_wm, mask_wm = item
img_raw, img_wm, mask_wm = img_raw.to(device), img_wm.to(
device), mask_wm.to(device)
mask, recon = model(img_wm)
# 加入discriminator
if args.gan_method:
# optimize D
dis_real, dis_recon = dis_forward(model_dis, img_raw,
recon.detach())
dis_wm = torch.sigmoid(model_dis(img_wm))
assert dis_recon.size() == dis_wm.size()
dis_fake = 0.5 * (dis_recon + dis_wm)
loss_disc = torch.mean(-1 * torch.log(1 - dis_fake) -
torch.log(dis_real))
loss_gp = net.calc_gradient_penalty(
model_dis, img_raw, recon.detach())
loss_d = loss_gp + loss_disc
# optimize G through D
dis_real, dis_recon = dis_forward(model_dis, img_raw,
recon)
# print('dis_real:', dis_real.size(), 'dis_recon:', dis_recon.size())
loss_g = 0.001 * torch.mean(-1 * torch.log(dis_recon))
loss_mask_reg = 0.1 * mask.clamp(0, 1).mean()
# loss_mask = 1000*util.exclusion_loss(mask)
try:
loss_mask = loss_bce(mask.clamp(0., 1.),
mask_wm.float().clamp(0., 1.))
except Exception:
import pdb
pdb.set_trace()
if not (mask >= 0. & mask <= 1.).all():
print('错误出在生成的mask')
if not (mask_wm >= 0. & mask_wm <= 1.).all():
print('错误出在gt水印')
loss_recon = loss(recon, img_raw)
loss_ssim = 1 - ssim._ssim(0.5 * (1 + img_raw), 0.5 *
(1 + recon), ssim_window, 11, 3,
True)
loss_weighted_recon = util.weighted_l1(recon, img_raw, mask)
loss_ = loss_recon + loss_mask + loss_ssim
if args.gan_method:
loss_ += loss_g
optimizer_dis.zero_grad()
loss_d.backward()
optimizer_dis.step()
optimizer.zero_grad()
loss_.backward()
optimizer.step()
epoch_loss += loss_.item()
step = i * batch_per_epoch + j
if j % 5 == 0:
writer.add_scalars(
'loss', {
'recon_l1': loss_recon.item(),
'ssim': loss_ssim.item(),
'exclusion': loss_mask.item(),
'mask_reg': loss_mask_reg.item()
}, step)
if j % 10 == 0:
print(
'Loss: %.3f (recon: %.3f \t ssim: %.3f \t mask: %.3f \t)'
% (loss_.item(), loss_recon.item(), loss_ssim.item(),
loss_mask.item()))
if args.gan_method:
print(
'disc: %.3f \t gp: %.3f \t gen: %.3f' %
(loss_disc.item(), loss_gp.item(), loss_g.item()))
# 记录mask和原图
if j % 50 == 0:
#import pdb; pdb.set_trace()
writer.add_images('images', [
torch.cat(3 * [mask[0].float().to(device)]),
torch.cat(
3 * [mask_wm[0].float().to(device).unsqueeze(0)]),
util.denormalize(img_wm[0]),
util.denormalize(recon[0]).clamp(0, 1)
],
global_step=step,
dataformats='CHW')
'''
writer.add_image('mask', mask[0], step)
writer.add_image('img', util.denormalize(img_wm[0]), step)
writer.add_image('recon_c', util.denormalize(recon_coarse[0]).clamp(0, 1), step)
writer.add_image('recon_f', util.denormalize(recon_fine[0]).clamp(0, 1), step)
writer.add_image('mask_gt', mask_wm[0], step)
'''
# 画各层的梯度分布图
if j % 100 == 0:
writer.add_figure('grad_flow',
util.plot_grad_flow_v2(
model.named_parameters()),
global_step=step)
if args.gan_method:
writer.add_figure('discriminator_grad',
util.plot_grad_flow_v2(
model_dis.named_parameters()),
global_step=step)
ckpt = {
'ckpt': model.state_dict(),
'optim': optimizer.state_dict()
}
torch.save(ckpt, os.path.join(args.save_dir, 'latest.pth'))
# 记录所有最好的epoch weight
if epoch_loss / (j + 1) < best_loss:
best_loss = epoch_loss / (j + 1)
shutil.copy(
os.path.join(args.save_dir, 'latest.pth'),
os.path.join(args.save_dir, 'epoch_' + str(i) + '.pth'))
except Exception as e:
ckpt = {'ckpt': model.state_dict(), 'optim': optimizer.state_dict()}
torch.save(ckpt, os.path.join(args.save_dir, 'latest.pth'))
print('Save temporary checkpoints to %s' % args.save_dir)
print(str(e), traceback.print_exc())
sys.exit(0)
print('Done training.')
shutil.copyfile(os.path.join(args.save_dir, 'latest.pth'),
os.path.join(args.save_dir, 'epoch_%d.pth' % (i + 1)))
def main():
parser = argparse.ArgumentParser(description='Train CAGAN')
# common
parser.add_argument('--n_iter', default=10000, type=int,
help='number of update')
parser.add_argument('--lr', default=0.0002, type=float,
help='learning rate of Adam')
parser.add_argument('--cuda', default=0,
help='0 indicates CPU')
parser.add_argument('--out', default='results',
help='log file')
parser.add_argument('--log_interval', default=20, type=int,
help='log inteval')
parser.add_argument('--ckpt_interval', default=50, type=int,
help='save interval')
parser.add_argument('--seed', default=42)
# model
parser.add_argument('--deconv', default='upconv',
help='deconv or upconv')
parser.add_argument('--relu', default='relu',
help='relu or relu6')
parser.add_argument('--bias', default=True,
help='use bias or not')
parser.add_argument('--init', default=None,
help='weight initialization')
# loss
parser.add_argument('--gamma_cycle', default=1.0, type=float,
help='coefficient for cycle loss')
parser.add_argument('--gamma_id', default=1.0, type=float,
help='coefficient for mask')
parser.add_argument('--norm', default=1, type=int,
help='selct norm type. Default is 1')
parser.add_argument('--cycle_norm', default=2, type=int,
help='selct norm type. Default is 2')
# dataset
parser.add_argument('--root', default='data',
help='root directory')
parser.add_argument('--base_root', default='images',
help='root directory to images')
parser.add_argument('--triplet', default='triplet.json',
help='triplet list')
args = parser.parse_args()
time = dt.now().strftime('%m%d_%H%M')
print('+++++ begin at {} +++++'.format(time))
for key, value in dict(args):
print('### {}: {}'.format(key, value))
print('+++++')
if args.seed is not None:
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
out = os.path.join(args.out, time)
if not os.path.isdir(out):
os.makedirs(out)
log_dir = os.path.join(out, 'tensorboard')
os.makedirs(log_dir)
ckpt_dir = os.path.join(out, 'checkpoints')
os.makedirs(ckpt_dir)
args.out = out
args.log_dir = log_dir
args.ckpt_dir = ckpt_dir
# logger
logger = getLogger()
f_h = FileHandler(os.path.join(out, 'log.txt'), 'a', 'utf-8')
f_h.setLevel(DEBUG)
s_h = StreamHandler()
s_h.setLevel(DEBUG)
logger.setLevel(DEBUG)
logger.addHandler(f_h)
logger.addHandler(s_h)
# tensorboard
writer = SummaryWriter(log_dir)
logger.debug("=====")
logger.debug(json.dumps(args.__dict__, indent=2))
logger.debug("=====")
kwargs = {'num_workers': 4, 'pin_memory': True} if args.cuda else {}
train_loader = torch.utils.data.DataLoader(
dataset.TripletDataset(
args.root,
os.path.join(args.root, args.triplet),
transform=transforms.Compose([
transforms.Scale((132, 100)),
transforms.RandomCrop((128, 96)),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.234, 0.225]),
transforms.ToTensor()])),
batch_size=args.batch_size, shuffle=True, **kwargs)
generator = net.Generator(args.deconv, args.relu, args.bias, args.norm)
discriminator = net.Discriminator(args.relu, args.bias)
opt_G = torch.optim.Adam(generator.parameters(), lr=args.lr)
opt_D = torch.optim.Adam(discriminator.parameters(), lr=args.lr)
if args.cuda:
generator = generator.cuda()
discriminator = discriminator.cuda()
start = dt.now()
logger.debug('===start {}'.format(start.strftime("%m/%d, %H:%M")))
for iter_ in range(args.n_iter):
# register params to tensorboard
if args.cuda:
generator = generator.cpu()
discriminator = discriminator.cpu()
for name, params in generator.named_children():
if params.requires_grad():
writer.add_histogram(tag='generator/' + name,
values=params.data.numpy(),
global_step=iter_)
for name, params in discriminator.named_children():
if params.requires_grad():
writer.add_histogram(tag='discriminator/' + name,
values=params.data.numpy(),
global_step=iter_)
if args.cuda:
generator = generator.cuda()
discriminator = discriminator.cuda()
list_of_tensor = train_loader.next()
# y_1: exchanged, y_2: cycle
loss_d, loss_g, mask_norm, cycle_loss, y_1, y_2 = train(
iter_, generator, discriminator, opt_G, opt_D, list_of_tensor,
args.cycle_norm, args.gamma_cycle, args.gamma_id, args.cuda)
writer.add_scalars(
main_tag='training',
tag_scalar_dict={
'discriminator/loss': loss_d,
'generator/loss': loss_d,
'generator/mask_norm': mask_norm,
'generator/cycle_loss': cycle_loss
}, global_step=iter_)
if iter_ % args.log_interval == 0:
msg = "Iter {} \tDis loss: {:.5f}\tGen loss: {:.5f}\tNorm: {:.5f}\n"
logger.debug(
msg.format(iter_, loss_d, loss_g, mask_norm + cycle_loss))
if iter_ % args.ckpt_interval == 0:
# save checkpoint and images used in this iteration
if args.cuda:
generator = generator.cpu()
discriminator = discriminator.cpu()
generator.eval()
discriminator.eval()
torch.save({
'iteration': iter_,
'gen_state': generator.state_dict(),
'dis_state': discriminator.state_dict(),
'opt_gen': opt_G.state_dict(),
'opt_dis': opt_D.state_dict()},
os.path.join(ckpt_dir, 'ckpt_iter_{}.pth'.format(iter_)))
# save images
if args.cuda:
y_1 = y_1.cpu()
y_2 = y_2.cpu()
writer.add_image(tag='input/human',
image_tensor=list_of_tensor[0],
global_step=iter_)
writer.add_image(tag='input/item',
image_tensor=list_of_tensor[1],
global_step=iter_)
writer.add_image(tag='input/want'
image_tensor=list_of_tensor[2],
global_step=iter_)
writer.add_image(tag='output/changed',
image_tensor=y_1.data,
global_step=iter_)
writer.add_image(tag='output/cycle',
image_tensor=y_2.data,
global_step=iter_)
if args.cuda:
generator = generator.cuda()
discriminator = discriminator.cuda()
# save checkpoint
torch.save({
'iteration': args.n_iter,
'gen_state': generator.state_dict(),
'dis_state': discriminator.state_dict(),
'opt_gen': opt_G.state_dict(),
'opt_dis': opt_D.state_dict()},
os.path.join(ckpt_dir, 'ckpt_iter_{}.pth'.format(args.n_iter)))
writer.export_scalars_to_json(os.path.join(log_dir + 'scalars.json'))
writer.close()
end = dt.now()
logger.debug('===end {}, {}[min]'.format(end.strftime(
'%m/%d, %H:%M'), (end - start).total_seconds() / 60.))
def main():
# This enables a ctr-C without triggering errors
import signal
signal.signal(signal.SIGINT, lambda x, y: sys.exit(0))
parser = argparse.ArgumentParser(description='GAN practice on MNIST')
parser.add_argument('--gpu',
'-g',
type=int,
default=-1,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--epoch',
'-e',
type=int,
default=100,
help='number of epochs to learn')
parser.add_argument('--batchsize',
'-b',
type=int,
default=100,
help='learning minibatch size')
parser.add_argument('--out',
'-o',
type=str,
default='model',
help='path to the output directory')
parser.add_argument('--dimz',
'-z',
type=int,
default=20,
help='dimention of encoded vector')
parser.add_argument('--seed',
type=int,
default=0,
help='Random seed of z at visualization stage')
parser.add_argument('--snapepoch',
'-s',
type=int,
default=10,
help='number of epochs to snapshot')
parser.add_argument('--load_gen_model',
type=str,
default='',
help='load generator model')
parser.add_argument('--load_dis_model',
type=str,
default='',
help='load generator model')
args = parser.parse_args()
if not os.path.exists(args.out):
os.makedirs(args.out)
print(args)
gen = net.Generator(784, args.dimz, 500)
dis = net.Discriminator(784, 500)
if args.load_gen_model != '':
chainer.serializers.load_npz(args.load_gen_model, gen)
print('Generator model loaded successfully!')
if args.load_dis_model != '':
chainer.serializers.load_npz(args.load_dis_model, dis)
print('Discriminator model loaded successfully!')
if args.gpu >= 0:
cuda.get_device(args.gpu).use()
gen.to_gpu()
dis.to_gpu()
print('GPU {}'.format(args.gpu))
xp = np if args.gpu < 0 else cuda.cupy
opt_gen = chainer.optimizers.Adam()
opt_dis = chainer.optimizers.Adam()
opt_gen.setup(gen)
opt_dis.setup(dis)
dataset = MnistDataset('./data')
# train, val = chainer.datasets.split_dataset_random(dataset, int(len(dataset) * 0.9))
train_iter = chainer.iterators.SerialIterator(dataset,
args.batchsize,
shuffle=True)
# val_iter = chainer.iterators.SerialIterator(val, args.batchsize, repeat=False, shuffle=False)
updater = GANUpdater(models=(gen, dis),
iterators={'main': train_iter},
optimizers={
'gen': opt_gen,
'dis': opt_dis
},
device=args.gpu,
params={
'batchsize': args.batchsize,
'n_latent': args.dimz
})
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
snapshot_interval = (args.snapepoch, 'epoch')
display_interval = (100, 'iteration')
trainer.extend(
extensions.snapshot(filename='snapshot_epoch_{.updater.epoch}.npz'),
trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(gen, 'gen{.updater.epoch}.npz'),
trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(dis, 'dis{.updater.epoch}.npz'),
trigger=snapshot_interval)
log_keys = ['epoch', 'iteration', 'gen/loss', 'dis/loss']
trainer.extend(
extensions.LogReport(keys=log_keys, trigger=display_interval))
trainer.extend(extensions.PrintReport(log_keys), trigger=display_interval)
trainer.extend(extensions.ProgressBar(update_interval=10))
trainer.extend(out_generated_image(gen, 10, 10, args.seed, args.out),
trigger=(1, 'epoch'))
trainer.run()
def gan_training(args, train):
# These iterators load the images with subprocesses running in parallel to
# the training/validation.
# Basically we define the batch size, number of processes which run together
# in each iteration; it also needs a 'train' object
if args.loaderjob:
train_iter = chainer.iterators.MultiprocessIterator(
train, args.batchsize, n_processes=args.loaderjob)
else:
train_iter = chainer.iterators.SerialIterator(train, args.batchsize)
# Prepare Flow GAN model, defined in net.py
gen = net.Generator(video_len=args.video_len)
dis = net.Discriminator()
if args.gpu >= 0:
cuda.get_device(args.gpu).use()
gen.to_gpu()
dis.to_gpu()
xp = np if args.gpu < 0 else cuda.cupy
# Setup optimizer to use to minimize the loss function
opt_gen = make_optimizer(gen, args)
opt_dis = make_optimizer(dis, args)
# Updater (updates parameters to train)
updater = GAN_Updater(models=(gen, dis),
iterator=train_iter,
optimizer={
'gen': opt_gen,
'dis': opt_dis
},
device=args.gpu)
# Trainer updates the params, including doing mini-batch loading, forward,
# backward computations, and executing update formula
trainer = training.Trainer(updater, (args.iteration, 'iteration'),
out=args.out)
snapshot_interval = (args.snapshot_interval), 'iteration'
visualize_interval = (args.visualize_interval), 'iteration'
log_interval = (args.log_interval), 'iteration'
# The Trainer class also invokes the extensions in decreasing order of priority
trainer.extend(extensions.LogReport(trigger=log_interval))
trainer.extend(
extensions.PlotReport(['gen/loss', 'dis/loss'],
trigger=log_interval,
file_name='plot.png'))
trainer.extend(extensions.PrintReport(
['epoch', 'iteration', 'gen/loss', 'dis/loss']),
trigger=log_interval)
trainer.extend(extensions.ProgressBar(update_interval=10))
trainer.extend(extensions.snapshot(), trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
gen, 'gen_iteration_{.updater.iteration}.npz'),
trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
dis, 'dis_iteration_{.updater.iteration}.npz'),
trigger=snapshot_interval)
trainer.extend(extension(gen, args), trigger=visualize_interval)
if args.adam_decay_iteration:
trainer.extend(extensions.ExponentialShift("alpha",
0.5,
optimizer=opt_gen),
trigger=(args.adam_decay_iteration, 'iteration'))
trainer.extend(extensions.ExponentialShift("alpha",
0.5,
optimizer=opt_dis),
trigger=(args.adam_decay_iteration, 'iteration'))
if args.resume:
chainer.serializers.load_npz(args.resume, trainer)
# Run the training
trainer.run()
from torch import Tensor as tensor
#pdb.set_trace()
gpuid = 0
lr_rate = 0.0002
alpha = 0.0005
beta = 0.01 * alpha
num_iter = 500000
optim_betas = (0.9, 0.999)
bs = 64
labelfake = Vb(torch.from_numpy(np.full((bs), 2, dtype=int))).cuda(gpuid)
labelpo = Vb(torch.from_numpy(np.full((bs), 1, dtype=int))).cuda(gpuid)
labelne = Vb(torch.from_numpy(np.full((bs), 0, dtype=int))).cuda(gpuid)
logging.basicConfig(filename='log/residualgan_v10.log', level=logging.INFO)
G1 = net.Generator1().cuda(gpuid)
G2 = net.Generator2().cuda(gpuid)
D = net.Discriminator().cuda(gpuid)
d_optimizer = optim.Adam(D.parameters(), lr=lr_rate, betas=optim_betas)
g1_optimizer = optim.Adam(G1.parameters(), lr=lr_rate, betas=optim_betas)
g2_optimizer = optim.Adam(G2.parameters(), lr=lr_rate, betas=optim_betas)
l1_crit = nn.L1Loss(size_average=False)
datalistpo = ld.getlist('../Eyeglasses_Positive.txt')
datalistne = ld.getlist('../Eyeglasses_Negative.txt')
iternow1 = 0
iternow2 = 0
for iter1 in xrange(num_iter):
D.zero_grad()
datapo, iternow1 = ld.load_data('../img_align_celeba_crop/',
'../Eyeglasses_Positive.txt',
datalistpo,
iternow1,
bs,
def main():
# Set random seem for reproducibility
manualSeed = 999
# manualSeed = random.randint(1,10000)
print("Random Seed: ",manualSeed)
random.seed(manualSeed)
torch.manual_seed(manualSeed)
# get device
device = torch.device("cuda:1" if (torch.cuda.is_available() and p.ngpu>0) else "cpu")
# get network
netG = net.Generator(p.ngpu).to(device)
netD = net.Discriminator(p.ngpu).to(device)
if (device.type=='cuda' and (p.ngpu > 1)):
netG = nn.DataParallel(netG,list(range(p.ngpu)))
netD = nn.DataParallel(netD,list(range(p.ngpu)))
#netG.apply(net.weights_init)
#netD.apply(net.weights_init)
print(netG)
print(netD)
# Loss function and optimizer
criterion = nn.BCELoss()
# create batch of latent vectors what we will use to visualize the progression of the generator
fixed_noise = torch.randn(p.batch_size,p.nz,1,1,device=device)
# Establish convention for real and fake labels during training
real_label = 1
fake_label = 0
# setup Adam optimiziers for both G and D
optimizerG = optim.Adam(netG.parameters(),lr=p.g_lr,betas=(p.beta1,0.999))
optimizerD = optim.Adam(netD.parameters(),lr=p.d_lr,betas=(p.beta1,0.999))
# start to train
img_list = []
G_losses = []
D_losses = []
iters = 0
for epoch in range(p.num_epochs):
for i,data in enumerate(dataset.dataloader,0):
# (1) Update D network: maximize log(D(x)) + log(1-D(G(z)))
netD.zero_grad()
# Format batch
real = data.to(device)
label = torch.full((data.size(0),),real_label,device=device)
# Forward pass real batch through D
output = netD(real).view(-1) # resize [batch_size,1,1,1] to [batch_size]
errD_real = criterion(output,label)
errD_real.backward()
D_x = output.mean().item()
#Generate batch of latent vectors
noise = torch.randn(data.size(0),p.nz,1,1,device=device)
#Generate batch of fake image
fake = netG(noise)
label.fill_(fake_label)
#Classify all fake batch with D
output = netD(fake.detach()).view(-1)
# Calculate D's loss on the all-fake batch
errD_fake = criterion(output,label)
errD_fake.backward()
D_G_z1 = output.mean().item()
errD = errD_real + errD_fake
#errD.backward()
optimizerD.step()
# (2) Update G network: maximize log(D(G(z)))
netG.zero_grad()
label.fill_(real_label) # fake labels are real for generator cost
output = netD(fake).view(-1)
errG = criterion(output,label)
# Calculate gradients for G
errG.backward()
D_G_z2 = output.mean().item()
optimizerG.step()
# output training stats
if i % 50 == 0:
print('[%d/%d][%d/%d]\tLoss_D: %.8f\tLoss_G: %.8f\tD(x): %.8f\tD(G(z)): %.8f / %.8f'
% (epoch, p.num_epochs, i, len(dataset.dataloader),
errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))
# Save Losses for plotting later
G_losses.append(errG.item())
D_losses.append(errD.item())
# Check how the generator is doing by saving G's output on fixed_noise
if (iters % 500 == 0) or ((epoch==p.num_epochs-1) and (i==len(dataset.dataloader)-1)):
with torch.no_grad():
fake = netG(fixed_noise).detach().cpu()
img_list.append(vutils.make_grid(fake,padding=2,normalize=True).numpy())
iters += 1
np.save(p.result_dir+"img_list.npy",img_list)
after_train(D_losses,G_losses)
show_after_train(img_list)
xp = enc.xp
is_AE = True
else:
gen = F.identity
xp = np
is_AE = False
print("Identity...")
## prepare networks for analysis
if args.output_analysis:
vgg = VGG16Layers() # for perceptual loss
vgg.to_gpu()
if is_AE:
enc_i = net.Encoder(args)
dec_i = net.Decoder(args)
dis = net.Discriminator(args)
dis_i = net.Discriminator(args)
if "enc_x" in args.load_models:
models = {
'enc_y': enc_i,
'dec_x': dec_i,
'dis_x': dis_i,
'dis_y': dis
}
else:
models = {
'enc_x': enc_i,
'dec_y': dec_i,
'dis_y': dis_i,
'dis_x': dis
}
def main():
args = easydict.EasyDict({
"dataroot": "/mnt/gold/users/s18150/mywork/pytorch/data/gan",
"save_dir": "./",
"prefix": "feature",
"workers": 8,
"batch_size": 128,
"image_size": 64,
"nc": 3,
"nz": 100,
"ngf": 64,
"ndf": 64,
"epochs": 50,
"lr": 0.0002,
"beta1": 0.5,
"gpu": 6,
"use_cuda": True,
"feature_matching": True,
"mini_batch": False
})
manualSeed = 999
random.seed(manualSeed)
torch.manual_seed(manualSeed)
device = torch.device(
'cuda:{}'.format(args.gpu) if args.use_cuda else 'cpu')
transform = transforms.Compose([
transforms.Resize(args.image_size),
transforms.CenterCrop(args.image_size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
dataset = dset.ImageFolder(root=args.dataroot, transform=transform)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers)
# Generator用のモデルのインスタンス作成
netG = net.Generator(args.nz, args.ngf, args.nc).to(device)
# Generator用のモデルの初期値を設定
netG.apply(net.weights_init)
# Discriminator用のモデルのインスタンス作成
netD = net.Discriminator(args.nc, args.ndf, device, args.batch_size,
args.mini_batch).to(device)
# Discriminator用のモデルの初期値を設定
netD.apply(net.weights_init)
# BCE Loss classのインスタンスを作成
criterionD = nn.BCELoss()
if args.feature_matching is True:
criterionG = nn.MSELoss(reduction='elementwise_mean')
else:
criterionG = nn.BCELoss()
# Generatorに入力するノイズをバッチごとに作成 (バッチ数は64)
# これはGeneratorの結果を描画するために使用する
fixed_noise = torch.randn(64, args.nz, 1, 1, device=device)
# 最適化関数のインスタンスを作成
optimizerD = optim.Adam(netD.parameters(),
lr=args.lr,
betas=(args.beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(),
lr=args.lr,
betas=(args.beta1, 0.999))
r = run.NNRun(netD, netG, optimizerD, optimizerG, criterionD, criterionG,
device, fixed_noise, args)
# 学習
r.train(dataloader)
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()
def gan_training(args, train):
# These iterators load the images with subprocesses running in parallel to
# the training/validation.
if args.loaderjob:
train_iter = chainer.iterators.MultiprocessIterator(
train, args.batchsize, n_processes=args.loaderjob)
else:
train_iter = chainer.iterators.SerialIterator(train, args.batchsize)
# Prepare Flow GAN model, defined in net.py
gen = net.Generator(video_len=args.video_len)
dis = net.Discriminator()
if args.gpu >= 0:
cuda.get_device(args.gpu).use()
gen.to_gpu()
dis.to_gpu()
xp = np if args.gpu < 0 else cuda.cupy
opt_gen = make_optimizer(gen, args)
opt_dis = make_optimizer(dis, args)
# Updater
updater = GAN_Updater(models=(gen, dis),
iterator=train_iter,
optimizer={
'gen': opt_gen,
'dis': opt_dis
},
device=args.gpu)
trainer = training.Trainer(updater, (args.iteration, 'iteration'),
out=args.out)
snapshot_interval = (args.snapshot_interval), 'iteration'
visualize_interval = (args.visualize_interval), 'iteration'
log_interval = (args.log_interval), 'iteration'
trainer.extend(extensions.LogReport(trigger=log_interval))
trainer.extend(
extensions.PlotReport(['gen/loss', 'dis/loss'],
trigger=log_interval,
file_name='plot.png'))
trainer.extend(extensions.PrintReport(
['epoch', 'iteration', 'gen/loss', 'dis/loss']),
trigger=log_interval)
trainer.extend(extensions.ProgressBar(update_interval=10))
trainer.extend(extensions.snapshot(), trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
gen, 'gen_iteration_{.updater.iteration}.npz'),
trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
dis, 'dis_iteration_{.updater.iteration}.npz'),
trigger=snapshot_interval)
trainer.extend(extension(gen, args), trigger=visualize_interval)
if args.adam_decay_iteration:
trainer.extend(extensions.ExponentialShift("alpha",
0.5,
optimizer=opt_gen),
trigger=(args.adam_decay_iteration, 'iteration'))
trainer.extend(extensions.ExponentialShift("alpha",
0.5,
optimizer=opt_dis),
trigger=(args.adam_decay_iteration, 'iteration'))
if args.resume:
chainer.serializers.load_npz(args.resume, trainer)
# Run the training
trainer.run()
