WindowsとMXNetでSRGAN(GPUなし)

はじめに

「Super-Resolution-Zoo」として各種学習済みモデルが公開されている。
github.com
今回はその中の「Super-Resolution-Zoo/SRGAN/SRGAN@leftthomas」内の「SRGAN_2x」を使わせて頂く。

環境

Windows7 64bit
Python 3.7.0

バージョン確認(pip freeze)

certifi==2018.8.24
chardet==3.0.4
graphviz==0.8.4
idna==2.6
mxnet==1.3.1b20181015
numpy==1.14.6
Pillow==5.3.0
requests==2.18.4
urllib3==1.22

実行

import numpy as np
import mxnet as mx
from mxnet import image

ctx = mx.cpu()

img = image.imread('sample.jpg')
img = img.astype(np.float32)/255
img = mx.nd.transpose(img, (2,0,1))
img = mx.nd.expand_dims(img, axis=0)

sym, arg_params, aux_params = mx.model.load_checkpoint('SRGAN_2x', 0)
model  mx.mod.Module(symbol=sym, label_names=None, context=ctx)
model.bind(for_training=False, data_shapes=[('data', img.shape)])
model.set_params(arg_params, aux_params)

from collections import namedtuple
Batch = namedtuple('Batch', ['data'])

model.forward(Batch([img]), is_train=False)
prob = model.get_outputs()[0].asnumpy()
prob = np.squeeze(prob)

from PIL import Image
prob = (prob.transpose(1,2,0)*255).astype(np.uint8)
img = Image.fromarray(prob)
img.save('output.jpg')

最後に

Python 3.7.0とmxnet-mklの組み合わせはうまくいかなかった