Refactored code for a Convolutional Autoencoder implemented with Chainer.

import argparse import numpy as np from chainer import Variable, FunctionSet, optimizers, cuda import chainer.functions as F import cv2 import random import cPickle as pickle import sys

class ConvolutionalAutoencoder(FunctionSet): def init(self, n_in, n_out, ksize, stride=1, pad=0, wscale=1, bias=0, nobias=False): super(ConvolutionalAutoencoder, self).init( encode=F.Convolution2D(n_in, n_out, ksize, stride=stride, pad=pad, wscale=wscale, bias=bias, nobias=nobias), decode=F.Convolution2D(n_out, n_in, ksize, stride=stride, pad=pad, wscale=wscale, bias=bias, nobias=nobias) )

def forward(self, x_data, y_data, train=True):
    x = Variable(x_data)
    t = Variable(x_data)
    if train:
        x = F.dropout(x)
    h = F.sigmoid(self.encode(x))
    y = F.sigmoid(self.decode(h))
    return F.mean_squared_error(y, t)

class Trainer(): def init(self, model, optimizer, loader, min_loss=None, max_epochs=10, logger=None): self.model = model self.optimizer = optimizer self.optimizer.setup(model.collect_parameters()) self.loader = loader self.min_loss = min_loss self.max_epochs = max_epochs self.logger = logger

def train(self, batchnum, batchsize):
    self.optimizer.zero_grads()
    x_data, y_data = self.loader(batchnum, batchsize)
    loss = self.model.forward(x_data, y_data, train=True)
    loss.backward()
    self.optimizer.update()
    return loss.data

def loop(self, N, batchsize):
    for epoch in xrange(self.max_epochs):
        for batchnum in xrange(N / batchsize):
            loss = float(cuda.to_cpu(self.train(batchnum, batchsize)))

            if not self.logger == None:
                logger(epoch+1, batchnum+1, loss)

            if not self.min_loss == None:
                if loss <= self.min_loss:
                    return

def load_image_list(path): tuples = [] for line in open(path): pair = line.strip().split() tuples.append((pair[0], np.int32(pair[1]))) return tuples

def compute(x_data, function): x = Variable(x_data) return function(x).data

def converter_generator(functions=[]): def converter(x_data): for function in functions: x_data = compute(x_data, function) return x_data

return converter

def read_image(path, insize, mean_image, center=False, flip=False): cropwidth = 256 - insize image = cv2.imread(path).transpose(2, 0, 1) if center: top = left = cropwidth / 2 else: top = random.randint(0, cropwidth - 1) left = random.randint(0, cropwidth - 1) bottom = insize + top right = insize + left

image  = image[[2, 1, 0], top:bottom, left:right].astype(np.float32)
image -= mean_image[:, top:bottom, left:right]
image /= 255
if flip and random.randint(0, 1) == 0:
    return image[:, :, ::-1]
else:
    return image

def loader_generator(insize, data, mean_image, converter, use_gpu=False): insize = 224 cropwidth = 256 - insize

perm = np.random.permutation(len(train_list))
def loader(batchnum, batchsize):
    x_batch = np.ndarray((batchsize, 3, insize, insize), dtype=np.float32)
    y_batch = np.ndarray((batchsize,), dtype=np.int32)

    for i in xrange(batchsize):
        path, label = data[perm[(batchnum * batchsize + i) % len(data)]]
        x_batch[i] = read_image(path, insize, mean_image, False, True)
        y_batch[i] = label

    if use_gpu:
        x_batch = cuda.to_gpu(x_batch)
        y_batch = cuda.to_gpu(y_batch)

    return converter(x_batch), y_batch

return loader

def logger(epoch, batchnum, loss): print 'Epoch {0:d} Batchnum {1:d} Loss={2:.5f}'.format(epoch, batchnum, loss)

if name == 'main': parser = argparse.ArgumentParser(description='Chainer example: Convolutional Autoencoder') parser.add_argument('--gpu', '-g', default=-1, type=int, help='GPU ID (negative value indicates CPU)') parser.add_argument('--batchsize', '-B', type=int, default=32, help='Learning minibatch size') parser.add_argument('--train', '-t', type=str, default="train.txt", help='Training image/label list') parser.add_argument('--val', '-v', type=str, default="val.txt", help='Validation image/label list') parser.add_argument('--mean', '-m', type=str, default='ilsvrc_2012_mean.npy', help='Image mean file') parser.add_argument('--conv1_1', type=str, default=None, help='Pickle file for conv1_1') args = parser.parse_args()

conv1_1 = ConvolutionalAutoencoder( 3, 64, 3, pad=1)

use_gpu = False
if args.gpu >= 0:
    use_gpu = True
    cuda.init(args.gpu)
    conv1_1.to_gpu()

insize = 224
train_list = load_image_list(args.train)
val_list = load_image_list(args.val)
mean_image = np.load(args.mean)
encoders = []
decoders = []

# Train layer 1
if args.conv1_1 == None:
    converter1_1 = converter_generator(encoders)
    loader1_1 = loader_generator(insize, train_list, mean_image, converter1_1, use_gpu=use_gpu)
    optimizer1_1 = optimizers.MomentumSGD(lr=0.01, momentum=0.9)
    trainer1_1 = Trainer(conv1_1, optimizer1_1, loader1_1, max_epochs=1, logger=logger)

    trainer1_1.loop(len(train_list), args.batchsize)

    f1_1 = open('pkl/conv1_1.pkl', 'wb')
    pickle.dump(conv1_1, f1_1)
    f1_1.close()
else:
    f1_1 = open(args.conv1_1, 'rb')
    conv1_1 = pickle.load(f1_1)
    f1_1.close()

encoders = encoders + [conv1_1.encode, F.sigmoid]
decoders = [conv1_1.decode, F.sigmoid] + decoders
reconstructor = converter_generator(encoders + decoders)

path, label = val_list[0]
x_data = np.ndarray((1, 3, insize, insize), dtype=np.float32)
x_data[0] = read_image(path, insize, mean_image)
if use_gpu:
    x_data = cuda.to_gpu(x_data)
y_data = cuda.to_cpu(reconstructor(x_data))
x_data = cuda.to_cpu(x_data)

origin = cv2.cvtColor(x_data[0].transpose(1, 2, 0) * 255, cv2.COLOR_RGB2BGR)
img1_1 = cv2.cvtColor(y_data[0].transpose(1, 2, 0) * 255, cv2.COLOR_RGB2BGR)
cv2.imwrite('img/origin.jpg', origin)
cv2.imwrite('img/img1_1.jpg', img1_1)