Pytorch上手使用

近期學習了另一個深度學習框架庫Pytorch，對學習進行一些總結，方便自己回顧。

Pytorch是torch的python版本，是由Facebook開源的神經網路框架。與Tensorflow的靜態計算圖不同，pytorch的計算圖是動態的，可以根據計算需要實時改變計算圖。

1 安裝

如果已經安裝了cuda8，則使用pip來安裝pytorch會十分簡單。若使用其他版本的cuda，則需要下載官方釋放出來對應的安裝包。具體安裝地址參見官網的首頁。
目前最新穩定版本為0.4.0。上個版本0.3.0的文件有中文版，見中文文件。

pip install torch torchvision  # for python2.7
 

pip3 install torch torchvision  # for python3

2 概述

理解pytorch的基礎主要從以下三個方面

• Numpy風格的Tensor操作。pytorch中tensor提供的API參考了Numpy的設計，因此熟悉Numpy的使用者基本上可以無縫理解，並建立和操作tensor，同時torch中的陣列和Numpy陣列物件可以無縫的對接。
• 變數自動求導。在一序列計算過程形成的計算圖中，參與的變數可以方便的計算自己對目標函式的梯度。這樣就可以方便的實現神經網路的後向傳播過程。
• 神經網路層與損失函式優化等高層封裝。網路層的封裝存在於torch.nn模組，損失函式由torch.nn.functional模組提供，優化函式由torch.optim模組提供。

因此下面的內容也主要圍繞這三個方面來介紹。第3節介紹張量的操作，第4節介紹自動求導，第5節介紹神經網路層等的封裝，第6，7節簡單介紹損失函式與優化方法。這三部分相對重要。後續的第8節介紹介紹資料集及torchvision，第9節介紹訓練過程可視的工具，第10節通過相對完整的示例程式碼展示pytorch中如何解決MNIST與CIFAR10的分類。

3 Tensor(張量)

Tensor是神經網路框架中重要的基礎資料型別，可以簡單理解為N維陣列的容器物件。tensor之間的通過運算進行連線，從而形成計算圖。

3.1 Tensor型別

Torch 定義了七種 CPU tensor 型別和八種 GPU tensor 型別:

Data type	CPU tensor	GPU tensor
32-bit floating point	torch.FloatTensor	torch.cuda.FloatTensor
64-bit floating point	torch.DoubleTensor	torch.cuda.DoubleTensor
16-bit floating point	torch.HalfTensor	torch.cuda.HalfTensor
8-bit integer (unsigned)	torch.ByteTensor	torch.cuda.ByteTensor
8-bit integer (signed)	torch.CharTensor	torch.cuda.CharTensor
16-bit integer (signed)	torch.ShortTensor	torch.cuda.ShortTensor
32-bit integer (signed)	torch.IntTensor	torch.cuda.IntTensor
64-bit integer (signed)	torch.LongTensor	torch.cuda.LongTensor

通常情況下使用Tensor類的建構函式返回的是FloatTensor型別物件，可通過在物件上呼叫cuda()返回一個新的cuda.FloatTensor型別的物件。

torch模組內提供了操作tensor的介面，而Tensor型別的物件上也設計了對應了介面。例如torch.add()與tensor.add()等價。需要注意的是這些介面都採用建立一個新物件返回的形式。如果想就地修改一個tensor物件，需要使用加字尾下劃線的方法。例如x.add_(y)，將修改x。Tensor類的構建函式支援從列表或ndarray等型別進行構建。預設tensor為FloatTensor。

下面的幾節簡單的描述重要的操作tensor的方法。

3.1 tensor的常見建立介面

方法名	說明
Tensor()	直接從引數構造一個的張量，引數支援list,numpy陣列
eye(row, column)	建立指定行數，列數的二維單位tensor
linspace(start,end,count)	在區間[s,e]上建立c個tensor
logspace(s,e,c)	在區間[10^s, 10^e]上建立c個tensor
ones(*size)	返回指定shape的張量，元素初始為1
zeros(*size)	返回指定shape的張量，元素初始為0
ones_like(t)	返回與t的shape相同的張量，且元素初始為1
zeros_like(t)	返回與t的shape相同的張量，且元素初始為0
arange(s,e,sep)	在區間[s,e)上以間隔sep生成一個序列張量

3.2 隨機取樣

方法名	說明
rand(*size)	在區間[0,1)返回一個均勻分佈的隨機數張量
uniform(s,e)	在指定區間[s,e]上生成一個均勻分佈的張量
randn(*size)	返回正態分佈N(0,1)取樣的隨機數張量
normal(means, std)	返回一個正態分佈N(means, std)

3.3 序列化

方法名	說明
save(obj, path)	張量物件的儲存，通過pickle進行
load(path)	從檔案中反序列化一個張量物件

3.4 數學操作

這些方法均為逐元素處理方法

方法名	說明
abs	絕對值
add	加法
addcdiv(t, v, t1, t2)	t1與t2的按元素除後，乘v加t
addcmul(t, v, t1, t2)	t1與t2的按元素乘後，乘v加t
ceil	向上取整
floor	向下取整
clamp(t, min, max)	將張量元素限制在指定區間
exp	指數
log	對數
pow	冪
mul	逐元素乘法
neg	取反
sigmoid
sign	取符號
sqrt	開根號
tanh

注:這些操作均建立新的tensor，如果需要就地操作，可以使用這些方法的下劃線版本，例如abs_。

3.5 歸約方法

方法名	說明
cumprod(t, axis)	在指定維度對t進行累積
cumsum	在指定維度對t進行累加
dist(a,b,p=2)	返回a,b之間的p階範數
mean	均值
median	中位數
std	標準差
var	方差
norm(t,p=2)	返回t的p階範數
prod(t)	返回t所有元素的積
sum(t)	返回t所有元素的和

3.6 比較方法

方法名	說明
eq	比較tensor是否相等，支援broadcast
equal	比較tensor是否有相同的shape與值
ge/le	大於/小於比較
gt/lt	大於等於/小於等於比較
max/min(t,axis)	返回最值，若指定axis，則額外返回下標
topk(t,k,axis)	在指定的axis維上取最高的K個值

3.7 其他操作

方法名	說明
cat(iterable, axis)	在指定的維度上拼接序列
chunk(tensor, c, axis)	在指定的維度上分割tensor
squeeze(input,dim)	將張量維度為1的dim進行壓縮，不指定dim則壓縮所有維度為1的維
unsqueeze(dim)	squeeze操作的逆操作
transpose(t)	計算矩陣的轉置換
cross(a, b, axis)	在指定維度上計算向量積
diag	返回對角線元素
hist(t, bins)	計算直方圖
trace	返回跡

3.8 矩陣操作

方法名	說明
dot(t1, t2)	計算張量的內積
mm(t1, t2)	計算矩陣乘法
mv(t1, v1)	計算矩陣與向量乘法
qr(t)	計算t的QR分解
svd(t)	計算t的SVD分解

3.9 tensor物件的方法

方法名	作用
size()	返回張量的shape屬性值
numel(input)	計算tensor的元素個數
view(*shape)	修改tensor的shape，與np.reshape類似，view返回的物件共享記憶體
resize	類似於view，但在size超出時會重新分配記憶體空間
item	若為單元素tensor，則返回pyton的scalar
from_numpy	從numpy資料填充
numpy	返回ndarray型別

3.10 tensor內部

tensor物件由兩部分組成，tensor的資訊與儲存，storage封裝了真正的data，可以由多個tensor共享。大多數操作只是修改tensor的資訊，而不修改storage部分。這樣達到效率與效能的提升。

3.11 使用pytorch進行線性迴歸

import torch
import torch.optim as optim
import matplotlib.pyplot as plt

def get_fake_data(batch_size=32):
    ''' y=x*2+3 '''
    x = torch.randn(batch_size, 1) * 20
    y = x * 2 + 3 + torch.randn(batch_size, 1)
    return x, y

x, y = get_fake_data()

class LinerRegress(torch.nn.Module):
    def __init__(self):
        super(LinerRegress, self).__init__()
        self.fc1 = torch.nn.Linear(1, 1)

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


net = LinerRegress()
loss_func = torch.nn.MSELoss()
optimzer = optim.SGD(net.parameters())

for i in range(40000):
    optimzer.zero_grad()

    out = net(x)
    loss = loss_func(out, y)
    loss.backward()

    optimzer.step()

w, b = [param.item() for param in net.parameters()]
print w, b  # 2.01146, 3.184525

# 顯示原始點與擬合直線
plt.scatter(x.squeeze().numpy(), y.squeeze().numpy())
plt.plot(x.squeeze().numpy(), (x*w + b).squeeze().numpy())
plt.show()

從這裡的程式碼可以發現，pytorch需要我們自己實現各輪更新，並手動呼叫反向傳播以及更新引數，此外也沒有提供評估及預測功能。相對於Keras這種高層的封裝，pytorch需要我們瞭解更多的低層細節。

4 自動求導

tensor物件通過一系列的運算可以組成動態圖，對於每個tensor物件，有下面幾個變數控制求導的屬性。

變數	作用
requirs_grad	預設為False，表示變數是否需要計算導數
grad_fn	變數的梯度函式
grad	變數對應的梯度

在0.3.0版本中，自動求導還需要藉助於Variable類來完成，在0.4.0版本中，Variable已經被廢除了，tensor自身即可完成這一過程。

import torch

x = torch.randn((4,4), requires_grad=True)
y = 2*x
z = y.sum()

print z.requires_grad  # True

z.backward()

print x.grad
'''
tensor([[ 2.,  2.,  2.,  2.],
        [ 2.,  2.,  2.,  2.],
        [ 2.,  2.,  2.,  2.],
        [ 2.,  2.,  2.,  2.]])
'''

5 建立神經網路

5.1 神經網路層

torch.nn模組提供了建立神經網路的基礎構件，這些層都繼承自Module類。下面我們簡單看下如何實現Liner層。

class Liner(torch.nn.Module):
    def __init__(self，in_features, out_features, bias=True):
        super(Liner, self).__init__()
        self.weight = torch.nn.Parameter(torch.randn(out_features, in_features))
        if bias:
            self.bias = torch.nn.Parameter(torch.randn(out_features))

    def forward(self, x):
        x = x.mm(self.weight)
        if self.bias:
            x = x + self.bias.expand_as(x)
        return x

下面表格中列出了比較重要的神經網路層元件。對應的在nn.functional模組中，提供這些層對應的函式實現。通常對於可訓練引數的層使用module，而對於不需要訓練引數的層如softmax這些，可以使用functional中的函式。

Layer對應的類	功能說明
Linear(in_dim, out_dim, bias=True)	提供了進行線性變換操作的功能
Dropout(p)	Dropout層，有2D,3D的型別
Conv2d(in_c, out_c, filter_size, stride, padding)	二維卷積層，類似的有Conv1d，Conv3d
ConvTranspose2d()
MaxPool2d(filter_size, stride, padding)	二維最大池化層
MaxUnpool2d(filter, stride, padding)	逆過程
AvgPool2d(filter_size, stride, padding)	二維平均池化層
FractionalMaxPool2d	分數最大池化
AdaptiveMaxPool2d([h,w])	自適應最大池化
AdaptiveAvgPool2d([h,w])	自自應平均池化
ZeroPad2d(padding_size)	零填充邊界
ConstantPad2d(padding_size,const)	常量填充邊界
ReplicationPad2d(ps)	複製填充邊界
BatchNorm1d()	對2維或3維小批量資料進行標準化操作
RNN(in_dim, hidden_dim, num_layers, activation, dropout, bidi, bias)	構建RNN層
RNNCell(in_dim, hidden_dim, bias, activation)	RNN單元
LSTM(in_dim, hidden_dim, num_layers, activation, dropout, bidi, bias)	構建LSTM層
LSTMCell(in_dim, hidden_dim, bias, activation)	LSTM單元
GRU(in_dim, hidden_dim, num_layers, activation, dropout, bidi, bias)	構建GRU層
GRUCell(in_dim, hidden_dim, bias, activation)	GRU單元

5.2 非線性啟用層

啟用層類名	作用
ReLU(inplace=False)	Relu啟用層
Sigmoid	Sigmoid啟用層
Tanh	Tanh啟用層
Softmax	Softmax啟用層
Softmax2d
LogSoftmax	LogSoftmax啟用層

5.3 容器型別

容器型別	功能
Module	神經網路模組的基類
Sequential	序列模型，類似於keras，用於構建序列型神經網路
ModuleList	用於儲存層，不接受輸入
Parameters(t)	模組的屬性，用於儲存其訓練引數
ParameterList	引數列表

下面的程式碼演示了使用容器型模組的方式。

# 方法1
model1 = nn.Sequential()
model.add_module('fc1', nn.Linear(3,4))
model.add_module('fc2', nn.Linear(4,2))
model.add_module('output', nn.Softmax(2))

# 方法2
model2 = nn.Sequential(
          nn.Conv2d(1,20,5),
          nn.ReLU(),
          nn.Conv2d(20,64,5),
          nn.ReLU()
        )
# 方法3        
model3 = nn.ModuleList([nn.Linear(3,4), nn.ReLU(), nn.Linear(4,2)])

5.4 其他層

容器型別	功能
Embedding(vocab_size, feature_dim)	詞向量層
Embeddingbag

5.5 模型的儲存

前面我們知道tensor可以通過save與load方法實現序列化與反序列化。由tensor組成的網路同樣也可以方便的儲存。不過通常沒有必要完全儲存網路模組物件，只需要儲存各層的權重資料即可，這些資料儲存在模組的state_dict字典中，因此只需要序列化這個詞典。

# 模型的儲存
torch.save(model.state_dict, 'path')
# 模型的載入
model.load_state_dict('path)

5.6 實現LeNet神經網路

torch.nn.Module提供了神經網路的基類，當實現神經網路時需要繼承自此模組，並在初始化函式中建立網路需要包含的層，並實現forward函式完成前向計算，網路的反向計算會由自動求導機制處理。

下面的示例程式碼建立了LeNet的卷積神經網路。通常將需要訓練的層寫在init函式中，將引數不需要訓練的層在forward方法裡呼叫對應的函式來實現相應的層。

import torch.nn as nn
import torch.nn.functional as F

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

6 損失函式與優化方法

6.1 損失函式

torch.nn模組中提供了許多損失函式類，這裡列出幾種相對常見的。

類名	功能
MSELoss	均方差損失
CrossEntropyLoss	交叉熵損失
NLLLoss	負對數似然損失
PoissonNLLLoss	帶泊松分佈的負對數似然損失

6.2 優化方法

由torch.optim模組提供支援

類名	功能
SGD(params, lr=0.1, momentum=0, dampening=0, weight_decay=0, nesterov=False)	隨機梯度下降法
Adam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)	Adam
RMSprop(params, lr=0.01, alpha=0.99, eps=1e-08, weight_decay=0, momentum=0, centered=False)	RMSprop
Adadelta(params, lr=1.0, rho=0.9, eps=1e-06, weight_decay=0)	Adadelta
Adagrad(params, lr=0.01, lr_decay=0, weight_decay=0, initial_accumulator_value=0)	Adagrad
lr_scheduler.ReduceLROnPlateau(optimizer, mode=’min’, factor=0.1, patience=10, verbose=False, threshold=0.0001, threshold_mode=’rel’, cooldown=0, min_lr=0, eps=1e-08)	學習率的控制

在神經網路的效能調優中，常見的作法是對不對層的網路設定不同的學習率。

class model(nn.Module):
    def __init__():
        super(model,self).__init__()
        self.base = Sequencial()
        # code for base sub module
        self.classifier = Sequencial()
        # code for classifier sub module

optim.SGD([
            {'params': model.base.parameters()},
            {'params': model.classifier.parameters(), 'lr': 1e-3}
            ], lr=1e-2, momentum=0.9)

6.3 引數初始化

良好的初始化可以讓模型快速收斂，有時甚至可以決定模型是否能訓練成功。Pytorch中的引數通常有預設的初始化策略，不需要我們自己指定，但框架仍然留有相應的介面供我們來調整初始化方法。

初始化方法	說明
xavier_uniform_
xavier_normal_
kaiming_uniform_

from torch.nn import init

# net的類定義
...

# 初始化各層權重
for name, params in net.named_parameters():
    init.xavier_normal(param[0])
    init.xavier_normal(param[1])

7 資料集與資料載入器

7.1 DataSet與DataLoader

torch.util.data模組提供了DataSet類用於描述一個數據集。定義自己的資料集需要繼承自DataSet類，且實現__getitem__()與__len__()方法。__getitem__方法返回指定索引處的tensor與其對應的label。

為了支援資料的批量及隨機化操作，可以使用data模組下的DataLoader型別來返回一個載入器：
DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, batch_sampler=None, num_workers=0)

7.2 torchvision簡介

torchvision是配合pytorch的獨立計算機視覺資料集的工具庫，下面介紹其中常用的資料集型別。

torchvision.datasets.ImageFolder(dir, transform, label_map,loader)
提供了從一個目錄初始化出來一個圖片資料集的便捷方法。
要求目錄下的圖片分類存放，每一類的圖片儲存在以類名為目錄名的目錄下，方法會將每個類名對映到唯一的數字上，如果你對數字有要求，可以用label_map來定義目錄名到數字的對映。

torchvision.datasets.DatasetFolder(dir,transform, label_map, loader, extensions)
提供了從一個目錄初始化一般資料集的便捷方法。目錄下的資料分類存放，每類資料儲存在class_xxx命名的目錄下。

此外torchvision.datasets下實現了常用的資料集，如CIFAR-10/100, ImageNet, COCO, MNIST, LSUN等。

除了資料集，torchvision的model模組提供了常見的模型實現，如Alex-Net, VGG，Inception, Resnet等。

7.3 torchvision提供的影象變換工具

torchvision的transforms模組提供了對PIL.Image物件和Tensor物件的常見操作。如果需要連續應用多個變換，可以使用Compose物件組裝多個變換。

轉換操作	說明
Scale	PIL圖片進行縮放
CenterCrop	PIL圖片從中心位置剪下
Pad	PIL圖片填充
ToTensor	PIL圖片轉換為Tensor且歸一化到[0,1]
Normalize	Tensor標準化
ToPILImage	將Tensor轉為PIL表示

import torchvision.tranforms as Trans
tranform = Trans.Compose([
                          T.Scale(28*28), 
                          T.ToTensor(), 
                          T.Normalize([0.5],[0.5])
                         ])

8 訓練過程視覺化

8.1 使用Tensorboard

通過使用第三方庫tensorboard_logger，將訓練過程中的資料儲存為日誌，然後便可以通過Tensorboard來檢視這些資料了。其功能支援相對有限，這裡不做過多介紹。

8.2 使用visdom

visdom是facebook開源的一個可視工具，可以用來完成pytorch訓練過程的視覺化。
安裝可以使用pip install visdom
啟動類似於tb，在命令列上執行：python -m visdom.server
服務啟動後可以使用瀏覽器開啟http://127.0.0.1:8097/即可看到主面板。

visdom的繪圖API類似於plot，通過API將繪圖資料傳送到基於tornado的web伺服器上並顯示在瀏覽器中。更詳細內容參見visdom的github主頁

9 GPU及並行支援

為了能在GPU上執行，Tensor與Module都需要轉換到cuda模式下。

import torch
import torchvision

t = torch.Tensor(3,4)
print t.is_cuda #False
t = t.cuda(0)
print t.is_cuda #True

net = torchvision.model.AlexNet()
net.cuda(0)

如果有多塊顯示卡，可以通過cuda(device_id)來將tensor分到不同的GPU上以達到負載的均衡。

另一種比較省事的做法是呼叫torch.set_default_tensor_type使程式預設使用某種cuda的tensor。或者使用torch.cuda.set_device(id)指定使用某個GPU。

10 示例：Pytorch實現CIFAR10與MNIST分類

關於cifar10與mnist資料集不再進行解釋了。這裡的Model類實現的二者的共同的任務，借鑑了keras的介面方式，Model類提供了train與evaluat方法，並沒有實現序列模型的新增方法以及predict方法。此外設定損失函式與優化函式時，也只是簡單的全部例項化出來，根據引數選擇其中的一個，這裡完全可以根據引數動態建立。

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms

class Model:
    def __init__(self, net, cost, optimist):
        self.net = net
        self.cost = self.create_cost(cost)
        self.optimizer = self.create_optimizer(optimist)
        pass

    def create_cost(self, cost):
        support_cost = {
            'CROSS_ENTROPY': nn.CrossEntropyLoss(),
            'MSE': nn.MSELoss()
        }

        return support_cost[cost]

    def create_optimizer(self, optimist, **rests):
        support_optim = {
            'SGD': optim.SGD(self.net.parameters(), lr=0.1, **rests),
            'ADAM': optim.Adam(self.net.parameters(), lr=0.01, **rests),
            'RMSP':optim.RMSprop(self.net.parameters(), lr=0.001, **rest)
        }

        return support_optim[optimist]

    def train(self, train_loader, epoches=3):
        for epoch in range(epoches):
            running_loss = 0.0
            for i, data in enumerate(train_loader, 0):
                inputs, labels = data

                self.optimizer.zero_grad()

                # forward + backward + optimize
                outputs = self.net(inputs)
                loss = self.cost(outputs, labels)
                loss.backward()
                self.optimizer.step()

                running_loss += loss.item()
                if i % 100 == 0:
                    print('[epoch %d, %.2f%%] loss: %.3f' %
                          (epoch + 1, (i + 1)*1./len(train_loader), running_loss / 100))
                    running_loss = 0.0

        print('Finished Training')

    def evaluate(self, test_loader):
        print('Evaluating ...')
        correct = 0
        total = 0
        with torch.no_grad():  # no grad when test and predict
            for data in test_loader:
                images, labels = data

                outputs = self.net(images)
                predicted = torch.argmax(outputs, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        print('Accuracy of the network on the test images: %d %%' % (100 * correct / total))


classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')


def cifar_load_data():
    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    trainset = torchvision.datasets.CIFAR10(root='./data', train=True,download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,shuffle=True, num_workers=2)

    testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
    testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                             shuffle=False, num_workers=2)

    return trainloader, testloader


class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

def mnist_load_data():
    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize([0,], [1,])])

    trainset = torchvision.datasets.MNIST(root='./data', train=True,
                                            download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=32,
                                              shuffle=True, num_workers=2)

    testset = torchvision.datasets.MNIST(root='./data', train=False,
                                           download=True, transform=transform)
    testloader = torch.utils.data.DataLoader(testset, batch_size=32,shuffle=True, num_workers=2)
    return trainloader, testloader


class MnistNet(torch.nn.Module):
    def __init__(self):
        super(MnistNet, self).__init__()
        self.fc1 = torch.nn.Linear(28*28, 512)
        self.fc2 = torch.nn.Linear(512, 512)
        self.fc3 = torch.nn.Linear(512, 10)

    def forward(self, x):
        x = x.view(-1, 28*28)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.softmax(self.fc3(x), dim=1)
        return x


if __name__ == '__main__':
    # train for mnist
    net = MnistNet()
    model = Model(net, 'CROSS_ENTROPY', 'RMSP')
    train_loader, test_loader = mnist_load_data()
    model.train(train_loader)
    model.evaluate(test_loader)

    # train for cifar
    net = LeNet()
    model = Model(net, 'CROSS_ENTROPY', 'RMSP')
    train_loader, test_loader = cifar_load_data()
    model.train(train_loader)
    model.evaluate(test_loader)