pytorch学习简单的nn和mnist分类

https://github.com/L1aoXingyu/code-of-learn-deep-learning-with-pytorch

参考以上网址的代码

import numpy as np
import torch
from torchvision.datasets import mnist
import matplotlib.pyplot as plt

from torch import nn
from torch.autograd import Variable

train_set = mnist.MNIST('./data',train=True,download=True)
test_set = mnist.MNIST('./data',train=False,download=True)
#a_data,a_label=train_set[0]
# plt.imshow(a_data)
# plt.show()
def data_tf(x):
    x = np.array(x,dtype='float32')/255
    x=(x-0.5)/0.5#转化为-1到1
    x=x.reshape((-1,))#拉成一行
    x=torch.from_numpy(x)
    return x

train_set=mnist.MNIST('./data',train=True,transform=data_tf,download=True)
test_set=mnist.MNIST('./data',train=False,transform=data_tf,download=True)

from torch.utils.data import DataLoader
train_data=DataLoader(train_set,batch_size=64,shuffle=True)
test_data=DataLoader(test_set,batch_size=128,shuffle=False)

a,a_label=next(iter(train_data))
net = nn.Sequential(nn.Linear(784,400),
                    nn.ReLU(),
                    nn.Linear(400,200),
                    nn.ReLU(),
                    nn.Linear(200,100),
                    nn.ReLU(),
                    nn.Linear(100,10))

#定义loss
criterion= nn.CrossEntropyLoss()
optimizer=torch.optim.SGD(net.parameters(),1e-1)
#开始训练
losses=[]
acces=[]
eval_losses=[]
eval_acces=[]
for e in range(20):
    train_loss=0
    train_acc=0
    net.train()
    for im,label in train_data:
        im = Variable(im)
        label = Variable(label)
        #前向传播
        out=net(im)
        loss=criterion(out,label)
        #反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        #记录误差
        train_loss+=loss.item()
        #计算分类准确率
        _,pred=out.max(1)
        num_correct=(pred==label).sum().item()
        acc=num_correct/im.shape[0]
        print(im.shape[0])#64
        train_acc+=acc
    losses.append(train_loss/len(train_data))
    acces.append(train_acc/len(train_data))#训练完所有数据的平均准确率
    eval_loss=0
    eval_acc=0
    net.eval()
    for im,label in test_data:
        im=Variable(im)
        label=Variable(label)
        out=net(im)
        loss=criterion(out,label)
        #记录误差
        eval_loss+=loss.item()
        #记录准确率
        _,pred=out.max(1)
        num_correct=(pred==label).sum().item()
        acc=num_correct/im.shape[0]
        eval_acc+=acc
    eval_losses.append(eval_loss/len(test_data))
    eval_acces.append(eval_acc/len(test_data))
    print('epoch:{},Train Loss:{:.6f},Train Acc:{:.6f},Eval Loss:{:.6f},Eval Acc:{:.6f}'
          .format(e,train_loss/len(train_data),train_acc/len(train_data),
                  eval_loss/len(test_data),eval_acc/len(test_data)))