def optimize(args):
""" Gatys et al. CVPR 2017
ref: Image Style Transfer Using Convolutional Neural Networks
"""
if args.cuda:
ctx = mx.gpu(0)
else:
ctx = mx.cpu(0)
# load the content and style target
content_image = utils.tensor_load_rgbimage(args.content_image,
ctx,
size=args.content_size,
keep_asp=True)
content_image = utils.subtract_imagenet_mean_preprocess_batch(
content_image)
style_image = utils.tensor_load_rgbimage(args.style_image,
ctx,
size=args.style_size)
style_image = utils.subtract_imagenet_mean_preprocess_batch(style_image)
# load the pre-trained vgg-16 and extract features
vgg = net.Vgg16()
utils.init_vgg_params(vgg, 'models', ctx=ctx)
# content feature
f_xc_c = vgg(content_image)[1]
# style feature
features_style = vgg(style_image)
gram_style = [net.gram_matrix(y) for y in features_style]
# output
output = Parameter('output', shape=content_image.shape)
output.initialize(ctx=ctx)
output.set_data(content_image)
# optimizer
trainer = gluon.Trainer([output], 'adam', {'learning_rate': args.lr})
mse_loss = gluon.loss.L2Loss()
# optimizing the images
for e in range(args.iters):
utils.imagenet_clamp_batch(output.data(), 0, 255)
# fix BN for pre-trained vgg
with autograd.record():
features_y = vgg(output.data())
content_loss = 2 * args.content_weight * mse_loss(
features_y[1], f_xc_c)
style_loss = 0.
for m in range(len(features_y)):
gram_y = net.gram_matrix(features_y[m])
gram_s = gram_style[m]
style_loss = style_loss + 2 * args.style_weight * mse_loss(
gram_y, gram_s)
total_loss = content_loss + style_loss
total_loss.backward()
trainer.step(1)
if (e + 1) % args.log_interval == 0:
print('loss:{:.2f}'.format(total_loss.asnumpy()[0]))
# save the image
output = utils.add_imagenet_mean_batch(output.data())
utils.tensor_save_bgrimage(output[0], args.output_image, args.cuda)
def optimize(args):
""" Gatys et al. CVPR 2017
ref: Image Style Transfer Using Convolutional Neural Networks
"""
if args.cuda:
ctx = mx.gpu(0)
else:
ctx = mx.cpu(0)
# load the content and style target
content_image = utils.tensor_load_rgbimage(args.content_image,ctx, size=args.content_size, keep_asp=True)
content_image = utils.subtract_imagenet_mean_preprocess_batch(content_image)
style_image = utils.tensor_load_rgbimage(args.style_image, ctx, size=args.style_size)
style_image = utils.subtract_imagenet_mean_preprocess_batch(style_image)
# load the pre-trained vgg-16 and extract features
vgg = net.Vgg16()
utils.init_vgg_params(vgg, 'models', ctx=ctx)
# content feature
f_xc_c = vgg(content_image)[1]
# style feature
features_style = vgg(style_image)
gram_style = [net.gram_matrix(y) for y in features_style]
# output
output = Parameter('output', shape=content_image.shape)
output.initialize(ctx=ctx)
output.set_data(content_image)
# optimizer
trainer = gluon.Trainer([output], 'adam',
{'learning_rate': args.lr})
mse_loss = gluon.loss.L2Loss()
# optimizing the images
for e in range(args.iters):
utils.imagenet_clamp_batch(output.data(), 0, 255)
# fix BN for pre-trained vgg
with autograd.record():
features_y = vgg(output.data())
content_loss = 2 * args.content_weight * mse_loss(features_y[1], f_xc_c)
style_loss = 0.
for m in range(len(features_y)):
gram_y = net.gram_matrix(features_y[m])
gram_s = gram_style[m]
style_loss = style_loss + 2 * args.style_weight * mse_loss(gram_y, gram_s)
total_loss = content_loss + style_loss
total_loss.backward()
trainer.step(1)
if (e + 1) % args.log_interval == 0:
print('loss:{:.2f}'.format(total_loss.asnumpy()[0]))
# save the image
output = utils.add_imagenet_mean_batch(output.data())
utils.tensor_save_bgrimage(output[0], args.output_image, args.cuda)
def train(args):
np.random.seed(args.seed)
if args.cuda:
ctx = mx.gpu(0)
else:
ctx = mx.cpu(0)
# dataloader
transform = utils.Compose([utils.Scale(args.image_size),
utils.CenterCrop(args.image_size),
utils.ToTensor(ctx),
])
train_dataset = data.ImageFolder(args.dataset, transform)
train_loader = gluon.data.DataLoader(train_dataset, batch_size=args.batch_size, last_batch='discard')
style_loader = utils.StyleLoader(args.style_folder, args.style_size, ctx=ctx)
print('len(style_loader):',style_loader.size())
# models
vgg = net.Vgg16()
utils.init_vgg_params(vgg, 'models', ctx=ctx)
style_model = net.Net(ngf=args.ngf)
style_model.initialize(init=mx.initializer.MSRAPrelu(), ctx=ctx)
if args.resume is not None:
print('Resuming, initializing using weight from {}.'.format(args.resume))
style_model.collect_params().load(args.resume, ctx=ctx)
print('style_model:',style_model)
# optimizer and loss
trainer = gluon.Trainer(style_model.collect_params(), 'adam',
{'learning_rate': args.lr})
mse_loss = gluon.loss.L2Loss()
for e in range(args.epochs):
agg_content_loss = 0.
agg_style_loss = 0.
count = 0
for batch_id, (x, _) in enumerate(train_loader):
n_batch = len(x)
count += n_batch
# prepare data
style_image = style_loader.get(batch_id)
style_v = utils.subtract_imagenet_mean_preprocess_batch(style_image.copy())
style_image = utils.preprocess_batch(style_image)
features_style = vgg(style_v)
gram_style = [net.gram_matrix(y) for y in features_style]
xc = utils.subtract_imagenet_mean_preprocess_batch(x.copy())
f_xc_c = vgg(xc)[1]
with autograd.record():
style_model.setTarget(style_image)
y = style_model(x)
y = utils.subtract_imagenet_mean_batch(y)
features_y = vgg(y)
content_loss = 2 * args.content_weight * mse_loss(features_y[1], f_xc_c)
style_loss = 0.
for m in range(len(features_y)):
gram_y = net.gram_matrix(features_y[m])
_, C, _ = gram_style[m].shape
gram_s = F.expand_dims(gram_style[m], 0).broadcast_to((args.batch_size, 1, C, C))
style_loss = style_loss + 2 * args.style_weight * mse_loss(gram_y, gram_s[:n_batch, :, :])
total_loss = content_loss + style_loss
total_loss.backward()
trainer.step(args.batch_size)
mx.nd.waitall()
agg_content_loss += content_loss[0]
agg_style_loss += style_loss[0]
if (batch_id + 1) % args.log_interval == 0:
mesg = "{}\tEpoch {}:\t[{}/{}]\tcontent: {:.3f}\tstyle: {:.3f}\ttotal: {:.3f}".format(
time.ctime(), e + 1, count, len(train_dataset),
agg_content_loss.asnumpy()[0] / (batch_id + 1),
agg_style_loss.asnumpy()[0] / (batch_id + 1),
(agg_content_loss + agg_style_loss).asnumpy()[0] / (batch_id + 1)
)
print(mesg)
if (batch_id + 1) % (4 * args.log_interval) == 0:
# save model
save_model_filename = "Epoch_" + str(e) + "iters_" + str(count) + "_" + str(time.ctime()).replace(' ', '_') + "_" + str(
args.content_weight) + "_" + str(args.style_weight) + ".params"
save_model_path = os.path.join(args.save_model_dir, save_model_filename)
style_model.collect_params().save(save_model_path)
print("\nCheckpoint, trained model saved at", save_model_path)
# save model
save_model_filename = "Final_epoch_" + str(args.epochs) + "_" + str(time.ctime()).replace(' ', '_') + "_" + str(
args.content_weight) + "_" + str(args.style_weight) + ".params"
save_model_path = os.path.join(args.save_model_dir, save_model_filename)
style_model.collect_params().save(save_model_path)
print("\nDone, trained model saved at", save_model_path)
def train():
if mx.context.num_gpus() > 0:
ctx = mx.gpu()
else:
raise RuntimeError('There is no GPU device!')
# loading configs
args = Options().parse()
cfg = Configs(args.config_path)
# set logging level
logging.basicConfig(level=logging.INFO)
# set random seed
np.random.seed(cfg.seed)
# build dataset and loader
content_dataset = ImageFolder(cfg.content_dataset, cfg.img_size, ctx=ctx)
style_dataset = StyleLoader(cfg.style_dataset, cfg.style_size, ctx=ctx)
content_loader = gluon.data.DataLoader(content_dataset, batch_size=cfg.batch_size, \
last_batch='discard')
vgg = Vgg16()
vgg._init_weights(fixed=True, pretrain_path=cfg.vgg_check_point, ctx=ctx)
style_model = Net(ngf=cfg.ngf)
if cfg.resume is not None:
print("Resuming from {} ...".format(cfg.resume))
style_model.collect_params().load(cfg.resume, ctx=ctx)
else:
style_model.initialize(mx.initializer.MSRAPrelu(), ctx=ctx)
print("Style model:")
print(style_model)
# build trainer
lr_sche = mx.lr_scheduler.FactorScheduler(
step=170000,
factor=0.1,
base_lr=cfg.base_lr
#warmup_begin_lr=cfg.base_lr/3.0,
#warmup_steps=300,
)
opt = mx.optimizer.Optimizer.create_optimizer('adam', lr_scheduler=lr_sche)
trainer = gluon.Trainer(style_model.collect_params(), optimizer=opt)
loss_fn = gluon.loss.L2Loss()
logging.info("Start training with total {} epoch".format(cfg.total_epoch))
iteration = 0
total_time = 0.0
num_batch = content_loader.__len__() * cfg.total_epoch
for epoch in range(cfg.total_epoch):
sum_content_loss = 0.0
sum_style_loss = 0.0
for batch_id, content_imgs in enumerate(content_loader):
iteration += 1
s = time.time()
style_image = style_dataset.get(batch_id)
style_vgg_input = subtract_imagenet_mean_preprocess_batch(
style_image.copy())
style_image = preprocess_batch(style_image)
style_features = vgg(style_vgg_input)
style_features = [
style_model.gram.gram_matrix(mx.nd, f) for f in style_features
]
content_vgg_input = subtract_imagenet_mean_preprocess_batch(
content_imgs.copy())
content_features = vgg(content_vgg_input)[1]
with autograd.record():
y = style_model(content_imgs, style_image)
y = subtract_imagenet_mean_batch(y)
y_features = vgg(y)
content_loss = 2 * cfg.content_weight * loss_fn(
y_features[1], content_features)
style_loss = 0.0
for m in range(len(y_features)):
gram_y = style_model.gram.gram_matrix(mx.nd, y_features[m])
_, C, _ = style_features[m].shape
gram_s = mx.nd.expand_dims(style_features[m],
0).broadcast_to((
gram_y.shape[0],
1,
C,
C,
))
style_loss = style_loss + 2 * cfg.style_weight * loss_fn(
gram_y, gram_s)
total_loss = content_loss + style_loss
total_loss.backward()
trainer.step(cfg.batch_size)
mx.nd.waitall()
e = time.time()
total_time += e - s
sum_content_loss += content_loss[0]
sum_style_loss += style_loss[0]
if iteration % cfg.log_interval == 0:
itera_sec = total_time / iteration
eta_str = str(
datetime.timedelta(seconds=int((num_batch - iteration) *
itera_sec)))
mesg = "{} Epoch [{}]:\t[{}/{}]\tTime:{:.2f}s\tETA:{}\tlr:{:.4f}\tcontent: {:.3f}\tstyle: {:.3f}\ttotal: {:.3f}".format(
time.strftime("%H:%M:%S",
time.localtime()), epoch + 1, batch_id + 1,
content_loader.__len__(), itera_sec, eta_str,
trainer.optimizer.learning_rate,
sum_content_loss.asnumpy()[0] / (batch_id + 1),
sum_style_loss.asnumpy()[0] / (batch_id + 1),
(sum_content_loss + sum_style_loss).asnumpy()[0] /
(batch_id + 1))
logging.info(mesg)
ctx.empty_cache()
save_model_filename = "Epoch_" + str(epoch + 1) + "_" + str(time.ctime()).replace(' ', '_') + \
"_" + str(cfg.content_weight) + "_" + str(cfg.style_weight) + ".params"
if not os.path.isdir(cfg.save_model_dir):
os.mkdir(cfg.save_model_dir)
save_model_path = os.path.join(cfg.save_model_dir, save_model_filename)
logging.info("Saving parameters to {}".format(save_model_path))
style_model.collect_params().save(save_model_path)
