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 test_check_gradient():
# test check gradient on a simple function
ctx = mx.cpu()
w = Parameter(name='w', shape=(2, 3))
w.initialize('zeros', ctx)
w.set_data(nd.array([[1., 2., 3], [-1., -3., 1.5]]))
def f():
return nd.sum(nd.square(w.data()))
assert check_gradient(f, [], w)
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)
class NormLinear(HybridBlock):
def __init__(self, num_features, num_classes, ctx, scale=15):
super(NormLinear, self).__init__()
self.num_classes = num_classes
self.scale = scale
with self.name_scope():
self.weight = Parameter('norm_weight', shape=(num_classes, num_features))
self.weight.initialize(init.Xavier(magnitude=2.24), ctx=ctx)
def hybrid_forward(self, F, x, weight):
feat = F.L2Normalization(x, mode='instance')
norm_weight = F.L2Normalization(weight, mode='instance')
cosine = F.FullyConnected(feat, norm_weight, no_bias=True, num_hidden=self.num_classes)
output = self.scale * cosine
return output # size=(Batch, Class)
