| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| import mxnet as mx |
| import numpy as np |
| import model_vgg19 as vgg |
| |
| class PretrainedInit(mx.init.Initializer): |
| def __init__(self, prefix, params, verbose=False): |
| self.prefix_len = len(prefix) + 1 |
| self.verbose = verbose |
| self.arg_params = {k : v for k, v in params.items() if k.startswith("arg:")} |
| self.aux_params = {k : v for k, v in params.items() if k.startswith("aux:")} |
| self.arg_names = set([k[4:] for k in self.arg_params.keys()]) |
| self.aux_names = set([k[4:] for k in self.aux_params.keys()]) |
| |
| def __call__(self, name, arr): |
| key = name[self.prefix_len:] |
| if key in self.arg_names: |
| if self.verbose: |
| print("Init %s" % name) |
| self.arg_params["arg:" + key].copyto(arr) |
| elif key in self.aux_params: |
| if self.verbose: |
| print("Init %s" % name) |
| self.aux_params["aux:" + key].copyto(arr) |
| else: |
| print("Unknown params: %s, init with 0" % name) |
| arr[:] = 0. |
| |
| |
| def style_gram_symbol(input_shape, style): |
| _, output_shapes, _ = style.infer_shape(**input_shape) |
| gram_list = [] |
| grad_scale = [] |
| for i in range(len(style.list_outputs())): |
| shape = output_shapes[i] |
| x = mx.sym.Reshape(style[i], shape=(int(shape[1]), int(np.prod(shape[2:])))) |
| # use fully connected to quickly do dot(x, x^T) |
| gram = mx.sym.FullyConnected(x, x, no_bias=True, num_hidden=shape[1]) |
| gram_list.append(gram) |
| grad_scale.append(np.prod(shape[1:]) * shape[1]) |
| return mx.sym.Group(gram_list), grad_scale |
| |
| |
| def get_loss(gram, content): |
| gram_loss = [] |
| for i in range(len(gram.list_outputs())): |
| gvar = mx.sym.Variable("target_gram_%d" % i) |
| gram_loss.append(mx.sym.sum(mx.sym.square(gvar - gram[i]))) |
| cvar = mx.sym.Variable("target_content") |
| content_loss = mx.sym.sum(mx.sym.square(cvar - content)) |
| return mx.sym.Group(gram_loss), content_loss |
| |
| def get_content_module(prefix, dshape, ctx, params): |
| sym = vgg.get_vgg_symbol(prefix, True) |
| init = PretrainedInit(prefix, params) |
| mod = mx.mod.Module(symbol=sym, |
| data_names=("%s_data" % prefix,), |
| label_names=None, |
| context=ctx) |
| mod.bind(data_shapes=[("%s_data" % prefix, dshape)], for_training=False) |
| mod.init_params(init) |
| return mod |
| |
| def get_style_module(prefix, dshape, ctx, params): |
| input_shape = {"%s_data" % prefix : dshape} |
| style, content = vgg.get_vgg_symbol(prefix) |
| gram, gscale = style_gram_symbol(input_shape, style) |
| init = PretrainedInit(prefix, params) |
| mod = mx.mod.Module(symbol=gram, |
| data_names=("%s_data" % prefix,), |
| label_names=None, |
| context=ctx) |
| mod.bind(data_shapes=[("%s_data" % prefix, dshape)], for_training=False) |
| mod.init_params(init) |
| return mod |
| |
| |
| def get_loss_module(prefix, dshape, ctx, params): |
| input_shape = {"%s_data" % prefix : dshape} |
| style, content = vgg.get_vgg_symbol(prefix) |
| gram, gscale = style_gram_symbol(input_shape, style) |
| style_loss, content_loss = get_loss(gram, content) |
| sym = mx.sym.Group([style_loss, content_loss]) |
| init = PretrainedInit(prefix, params) |
| gram_size = len(gram.list_outputs()) |
| mod = mx.mod.Module(symbol=sym, |
| data_names=("%s_data" % prefix,), |
| label_names=None, |
| context=ctx) |
| mod.bind(data_shapes=[("%s_data" % prefix, dshape)], |
| for_training=True, inputs_need_grad=True) |
| mod.init_params(init) |
| return mod, gscale |
| |
| |
| |
| if __name__ == "__main__": |
| from data_processing import PreprocessContentImage, PreprocessStyleImage |
| from data_processing import PostprocessImage, SaveImage |
| vgg_params = mx.nd.load("./model/vgg19.params") |
| style_weight = 2 |
| content_weight = 10 |
| long_edge = 384 |
| content_np = PreprocessContentImage("./input/IMG_4343.jpg", long_edge) |
| style_np = PreprocessStyleImage("./input/starry_night.jpg", shape=content_np.shape) |
| dshape = content_np.shape |
| ctx = mx.gpu() |
| # style |
| style_mod = get_style_module("style", dshape, ctx, vgg_params) |
| style_mod.forward(mx.io.DataBatch([mx.nd.array(style_np)], [0]), is_train=False) |
| style_array = [arr.copyto(mx.cpu()) for arr in style_mod.get_outputs()] |
| del style_mod |
| # content |
| content_mod = get_content_module("content", dshape, ctx, vgg_params) |
| content_mod.forward(mx.io.DataBatch([mx.nd.array(content_np)], [0]), is_train=False) |
| content_array = content_mod.get_outputs()[0].copyto(mx.cpu()) |
| del content_mod |
| # loss |
| mod, gscale = get_loss_module("loss", dshape, ctx, vgg_params) |
| extra_args = {"target_gram_%d" % i : style_array[i] for i in range(len(style_array))} |
| extra_args["target_content"] = content_array |
| mod.set_params(extra_args, {}, True, True) |
| grad_array = [] |
| for i in range(len(style_array)): |
| grad_array.append(mx.nd.ones((1,), ctx) * (float(style_weight) / gscale[i])) |
| grad_array.append(mx.nd.ones((1,), ctx) * (float(content_weight))) |
| # train |
| img = mx.nd.zeros(content_np.shape, ctx=ctx) |
| img[:] = mx.rnd.uniform(-0.1, 0.1, img.shape) |
| lr = mx.lr_scheduler.FactorScheduler(step=80, factor=.9) |
| optimizer = mx.optimizer.SGD( |
| learning_rate = 0.001, |
| wd = 0.0005, |
| momentum=0.9, |
| lr_scheduler = lr) |
| optim_state = optimizer.create_state(0, img) |
| |
| old_img = img.copyto(ctx) |
| clip_norm = 1 * np.prod(img.shape) |
| |
| import logging |
| for e in range(800): |
| mod.forward(mx.io.DataBatch([img], [0]), is_train=True) |
| mod.backward(grad_array) |
| data_grad = mod.get_input_grads()[0] |
| gnorm = mx.nd.norm(data_grad).asscalar() |
| if gnorm > clip_norm: |
| print("Data Grad: ", gnorm / clip_norm) |
| data_grad[:] *= clip_norm / gnorm |
| |
| optimizer.update(0, img, data_grad, optim_state) |
| new_img = img |
| eps = (mx.nd.norm(old_img - new_img) / mx.nd.norm(new_img)).asscalar() |
| old_img = new_img.copyto(ctx) |
| logging.info('epoch %d, relative change %f', e, eps) |
| if (e+1) % 50 == 0: |
| SaveImage(new_img.asnumpy(), 'output/tmp_'+str(e+1)+'.jpg') |
| |
| SaveImage(new_img.asnumpy(), "./output/out.jpg") |
| |
