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))))