pytorch minist全网络手写数字识别

#!/usr/bin/python
# -*- coding: UTF-8 -*-

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from torch.autograd import Variable
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np

class Batch_Net(torch.nn.Module):
    def __init__(self, in_dim, n_hidden_1, n_hidden_2, out_dim):
        super(Batch_Net, self).__init__() //固定格式
        self.layer1 = torch.nn.Sequential( //Sequential为顺序进行参数中的操作
            torch.nn.Linear(in_dim, n_hidden_1),  //第一层全连接层
            torch.nn.BatchNorm1d(n_hidden_1), torch.nn.ReLU(True) //将输出非线型化，并添加激活函数
        )
        self.layer2 = torch.nn.Sequential(
            torch.nn.Linear(n_hidden_1, n_hidden_2), //第二层全连接层
            torch.nn.BatchNorm1d(n_hidden_2), torch.nn.ReLU(True)
        )
        self.layer3 = torch.nn.Sequential(torch.nn.Linear(n_hidden_2, out_dim)) //输出层，输出层不能非线形化

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        return x

batch_size = 32  //每次训练以32张图片为一组，需要总数量整除
learning_rate = 1e-2 //学习效率
num_epoches = 20 //每次训练的批次
data_tf = transforms.Compose( #Compose 把多个步骤合到一起
    [transforms.ToTensor(),
     transforms.Normalize([0.5], [0.5])] #将图片转化到-１~１之间
)

train_dataset = datasets.MNIST( //下载训练集，由于我已经下载了，所以download为False
    root='./data', train=True, transform=data_tf, download=False
)
test_dataset = datasets.MNIST(root='./data', train=False, transform=data_tf, download=False)//训练集

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)//加载训练集  #ｓｈｕｆｆｌｅ表示每次迭代的时候是否将数据打乱
test_loader  = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)//加载测试集

model = Batch_Net(28 * 28, 300, 100, 10)//图片为28＊28的灰度图，因为数字为0－9所以输出层为10维
if(torch.cuda.is_available()):
    model = model.cuda()

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

for epoch in range(num_epoches):
    print('epoch {}'.format(epoch + 1))
    #training----------------------
    train_loss = 0.
    train_acc = 0.
    for img, label in train_loader:
        img = Variable(img)
        img = img.view(batch_size, 28 * 28) //将图片序列化成二维矩阵
        label = Variable(label)
        if(torch.cuda.is_available()):
            img = img.cuda()
            label = label.cuda()
        out = model(img)
        loss = criterion(out, label)
        train_loss += loss.item()
        pred = torch.max(out, 1)[1]
        train_correct = (pred.data.cpu() == label.data.cpu()).sum()
        train_acc += train_correct.item()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print('Train Loss: {:.6f}, Acc: {:.6f}'.format(train_loss / (len(train_loader)), train_acc / (len(train_loader))))