mnist_ae_gpu.py

import numpy as np
import chainer.functions as F
import chainer.links as L
from chainer import Variable,optimizers,Chain
import data
import cPickle as pickle


class Model(Chain):
    def __init__(self):
        super(Model, self).__init__(
            l1=L.Linear(784, 400),
            l1b=L.Linear(400, 784),
            l2=L.Linear(400, 100),
            l2b=L.Linear(100, 400),
            l3=L.Linear(100, 49),
            l3b=L.Linear(49, 100),
            l4=L.Linear(49, 16),
            l4b=L.Linear(16, 49),
            l5=L.Linear(16, 3),
            l5b=L.Linear(3, 16),
            )
    def __call__(self, x,train=True):
        h = F.dropout(F.relu(self.l1(x)),train=train)
        h = F.dropout(self.l1b(h),train=train)
        return h
    def l2_f(self, x,train=True):
        x = F.dropout(F.relu(self.l1(x)),train=False)
        h = F.dropout(F.relu(self.l2(x)),train=train)
        h = F.dropout(self.l2b(h),train=train)
        return h,F.mean_squared_error(x,h)
    def l3_f(self, x,train=True):
        x = F.dropout(F.relu(self.l1(x)),train=False)
        x = F.dropout(F.relu(self.l2(x)),train=False)
        h = F.dropout(F.relu(self.l3(x)),train=train)
        h = F.dropout(self.l3b(h),train=train)
        return h,F.mean_squared_error(x,h)
    def l4_f(self, x,train=True):
        x = F.dropout(F.relu(self.l1(x)),train=False)
        x = F.dropout(F.relu(self.l2(x)),train=False)
        x = F.dropout(F.relu(self.l3(x)),train=False)
        h = F.dropout(F.relu(self.l4(x)),train=train)
        h = F.dropout(self.l4b(h),train=train)
        return h,F.mean_squared_error(x,h)
    def l5_f(self, x,train=True):
        x = F.dropout(F.relu(self.l1(x)),train=False)
        x = F.dropout(F.relu(self.l2(x)),train=False)
        x = F.dropout(F.relu(self.l3(x)),train=False)
        x = F.dropout(F.relu(self.l4(x)),train=False)
        h = F.dropout(F.relu(self.l5(x)),train=train)
        h = F.dropout(self.l5b(h),train=train)
        return h,F.mean_squared_error(x,h)
    def dump(self):
        pickle.dump(self.l1,open('l1.pkl', 'w'))
        pickle.dump(self.l2,open('l2.pkl', 'w'))
        pickle.dump(self.l3,open('l3.pkl', 'w'))
        pickle.dump(self.l4,open('l4.pkl', 'w'))
        pickle.dump(self.l5,open('l5.pkl', 'w'))


from chainer import cuda
gpu=0
cuda.get_device(gpu).use()
xp = cuda.cupy
mnist = data.load_mnist_data()
x_all = mnist['data'].astype(xp.float32) / 255
y_all = mnist['target'].astype(xp.int32)
x_train, x_test = np.split(x_all, [60000])
y_train, y_test = np.split(y_all, [60000])


model = Model()
model.to_gpu(gpu)
optimizer = optimizers.MomentumSGD(lr=0.01, momentum=0.9)
optimizer.setup(model)



batchsize = 100
datasize = 60000
epochs = 500

import matplotlib.pyplot as plt

model.dump()

for j in range(epochs):
    indexes = np.random.permutation(datasize)
    for i in range(0, datasize, batchsize):
        x = Variable(xp.asarray(x_train[indexes[i : i + batchsize]]))
        t = Variable(xp.asarray(y_train[indexes[i : i + batchsize]]))
        model.zerograds()
        img = model(x)
        loss = F.mean_squared_error(img, x)
        loss.backward()
        optimizer.update()

    x = Variable(xp.asarray(x_test))
    img = model(x,train=False)
    loss = F.mean_squared_error(img, x)
    print "layer1:",j,loss.data

model.dump()
optimizer.setup(model)
for j in range(epochs):
    indexes = np.random.permutation(datasize)
    for i in range(0, datasize, batchsize):
        x = Variable(xp.asarray(x_train[indexes[i : i + batchsize]]))
        t = Variable(xp.asarray(y_train[indexes[i : i + batchsize]]))
        model.zerograds()
        img,loss = model.l2_f(x)
        #loss = F.mean_squared_error(img, x)
        loss.backward()
        optimizer.update()

    x = Variable(xp.asarray(x_test))
    img,loss = model.l2_f(x,train=False)
    print "layer2:",j,loss.data

model.dump()
optimizer.setup(model)

for j in range(epochs):
    indexes = np.random.permutation(datasize)
    for i in range(0, datasize, batchsize):
        x = Variable(xp.asarray(x_train[indexes[i : i + batchsize]]))
        t = Variable(xp.asarray(y_train[indexes[i : i + batchsize]]))
        model.zerograds()
        img,loss = model.l3_f(x)
        loss.backward()
        optimizer.update()

    x = Variable(xp.asarray(x_test))
    img,loss = model.l3_f(x,train=False)
    print "layer3:",j,loss.data

model.dump()
optimizer.setup(model)

for j in range(epochs):
    indexes = np.random.permutation(datasize)
    for i in range(0, datasize, batchsize):
        x = Variable(xp.asarray(x_train[indexes[i : i + batchsize]]))
        t = Variable(xp.asarray(y_train[indexes[i : i + batchsize]]))
        model.zerograds()
        img,loss = model.l4_f(x)
        loss.backward()
        optimizer.update()

    x = Variable(xp.asarray(x_test))
    img,loss = model.l4_f(x,train=False)
    print "layer4:",j,loss.data

model.dump()
optimizer.setup(model)

for j in range(epochs):
    indexes = np.random.permutation(datasize)
    for i in range(0, datasize, batchsize):
        x = Variable(xp.asarray(x_train[indexes[i : i + batchsize]]))
        t = Variable(xp.asarray(y_train[indexes[i : i + batchsize]]))
        model.zerograds()
        img,loss = model.l5_f(x)
        loss.backward()
        optimizer.update()

    x = Variable(xp.asarray(x_test))
    img,loss = model.l5_f(x,train=False)
    print "layer5:",j,loss.data
    #img = img.data[0].reshape(4,4)
    #plt.imshow(img)
    #plt.savefig("l5_encoded_%d.png"%j)

model.dump()