In [ ]:
# Block NULL: Readme

# Note by lzh: Always run Block A and Block B, and run Block C and / or Block D if necessary.
# Remeber to restart before conducting each experiment.

In [ ]:
# Block A: Set up matplotlib and jupyter notebook

# Note by lzh: use %matplotlib inline to force plots to be static, use %matplotlib notebook
# to permit interactive plots
%matplotlib notebook

In [ ]:
# Block B': Import packages and definition of NN

import torch
import torchvision
from torch.utils.data import DataLoader
from torchvision import transforms
from torch.autograd import Variable
from torch import optim
import torch.nn as nn
import torch.nn.functional as F

import numpy as np
from random import random
import math
from math import sqrt
from math import sin, pi

import datetime

import os

global_cuda_available = True

if global_cuda_available:
    os.environ["CUDA_VISIBLE_DEVICES"] = "1"

from matplotlib import pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D

class block(nn.Module):
    def __init__(self,inplanes,outplanes):
        super(block,self).__init__()
        self.fc = nn.Linear(inplanes,outplanes)

    def forward(self,x):
        x = F.relu(self.fc(x))
        return x

class TestNet(nn.Module):
    def __init__(self):
        super(TestNet, self).__init__()
        layer = [block(2,50)]
        for i in range(1):
            layer.append(block(50,50))

        self.fc = nn.Linear(50,1)
        layer.append(self.fc)
        self.forward_prop=nn.Sequential(*layer)


    def forward(self, x):
        return self.forward_prop(x)

f = lambda x: math.cos(x)
g = lambda x: math.sin(x)
test_func = lambda x: 10.0 * sin(20.0 * x)
n = 20
data_generate = [i / n * math.pi / 2.0 for i in range(n)]
dataset_x = Variable(torch.Tensor(list(map(lambda x:[f(x), g(x)], data_generate))))
dataset_y = Variable(torch.Tensor(list(map(test_func, data_generate))))

Net = TestNet()

if global_cuda_available:
    Net.cuda()

print(dataset_x)
print(dataset_y)

learning_rate = 1e-4

criterian = nn.MSELoss()
# optimizer = optim.Adam(Net.parameters(),lr=learning_rate)
optimizer = optim.SGD(Net.parameters(), lr = 0.001, momentum = 0.9)

min_los = float('inf')

In [ ]:
# Block C: Do regression with NN and save the model

model_filename = "Model-" + datetime.datetime.now().strftime("%Y%m%d%H%M%S%f") + ".mdlpkl"

print("Filename of the model is", model_filename)

iter_time = 50000
print_factor = 1000
for i in range(iter_time):
    
    optimizer.zero_grad()
    
    if global_cuda_available:
        output = Net(dataset_x.cuda())
    else:
        output = Net(dataset_x)

    if global_cuda_available:
        loss = criterian(output, dataset_y.cuda())
    else:
        loss = criterian(output, dataset_y)
    loss.backward()
    optimizer.step()

    running_loss = 0.
    running_acc = 0.
    running_loss += loss.data[0]
    if running_loss < min_los:
        min_los = running_loss
        if i % print_factor == 0:
            torch.save(Net, model_filename)
        output_r = output

    if i % print_factor == 0:
        print("[%d/%d] Loss: %.20f" % (i + 1, iter_time, running_loss))

print(min_los)
print(output_r, dataset_y)

In [ ]:
# Block D': Draw figures with matplotlib

# Note by lzh: If you want to load 2dmodel_.pkl, uncomment the following sentence
#     for 2dmodel_.pkl is generated from a different mechanism
# from test2D import f, n, TestNet, test_func,block

testnet = torch.load("Model-20170708235733720778.pkl")
testnet.eval()

def map_func(func):
    return lambda x:list(map(lambda x:func([[x[0],x[1]]]),x))

n1 = 35
grid_start = -0.1
grid_end = 1.1
grid_range = grid_end - grid_start

fig = plt.figure()
ax = Axes3D(fig)
X = np.arange(grid_start, grid_end, 1/n1)
Y = np.arange(grid_start, grid_end, 1/n1)
X, Y = np.meshgrid(X, Y)
data_test = np.column_stack([X.reshape(-1), Y.reshape(-1)]).astype(np.float32)

if global_cuda_available:
    data_plot = testnet.forward(Variable(torch.from_numpy(data_test)).cuda()).cpu().data.numpy().reshape(-1)
else:
    data_plot = testnet.forward(Variable(torch.from_numpy(data_test))).data.numpy().reshape(-1)

ax.plot_surface(X, Y, data_plot.reshape(X.shape), rstride=1, cstride=1, cmap='rainbow')

ax.scatter(dataset_x.data.numpy()[::, 0], dataset_x.data.numpy()[::, 1])
ax.scatter(dataset_x.data.numpy()[::, 0], dataset_x.data.numpy()[::, 1], dataset_y.data.numpy()[::], alpha=0.4, marker='+')

plt.show()

In [ ]: