pytorch-mnist神經網絡訓練

在net.py里面構造網絡，網絡的結構為輸入為28*28,第一層隱藏層的輸出為300, 第二層輸出的輸出為100, 最后一層的輸出層為10,  

net.py 

import torch
from torch import nn

class Batch_Net(nn.Module):
    def __init__(self, in_dim, n_hidden_1, n_hidden_2, out_dim):
        super(Batch_Net, self).__init__()
        self.layer_1 = nn.Sequential(nn.Linear(in_dim, n_hidden_1), nn.BatchNorm1d(n_hidden_1), nn.ReLU(True))
        self.layer_2 = nn.Sequential(nn.Linear(n_hidden_1, n_hidden_2), nn.BatchNorm1d(n_hidden_2), nn.ReLU(True))
        self.output = nn.Sequential(nn.Linear(n_hidden_2, out_dim))

    def forward(self, x):
        x = self.layer_1(x)
        x = self.layer_2(x)
        x = self.output(x)
        return x

main.py 進行網絡的訓練 

import torch
from torch import nn, optim
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


import net


batch_size = 128  # 每一個batch_size的大小
learning_rate = 1e-2 # 學習率的大小
num_epoches = 20  # 迭代的epoch值
 # 表示data將數據變成0, 1之間，0.5, 0.5表示減去均值處以標准差
data_tf = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])  # 表示均值和標准差
# 獲得訓練集的數據
train_dataset = datasets.MNIST(root='./data', train=True, transform=data_tf, download=True)
# 獲得測試集的數據
test_dataset = datasets.MNIST(root='./data', train=False, transform=data_tf, download=True)
# 獲得訓練集的可迭代隊列
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
# 獲得測試集的可迭代隊列
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
# 構造模型的網絡
model = net.Batch_Net(28*28, 300, 100, 10)
if torch.cuda.is_available(): # 如果有cuda就將模型放在GPU上
    model.cuda()

criterion = nn.CrossEntropyLoss() # 構造交叉損失函數
optimizer = optim.SGD(model.parameters(), lr=learning_rate) # 構造模型的優化器

for epoch in range(num_epoches): # 迭代的epoch 
    train_loss = 0 # 訓練的損失值
    test_loss = 0 # 測試的損失值
    eval_acc = 0 # 測試集的准確率
    for data in train_loader:  # 獲得一個batch的樣本 
        img, label = data # 獲得圖片和標簽
        img = img.view(img.size(0), -1) # 將圖片進行img的轉換
        if torch.cuda.is_available(): # 如果存在torch 
            img = Variable(img).cuda() # 將圖片放在torch上
            label = Variable(label).cuda() # 將標簽放在torch上 
        else:
            img = Variable(img)  # 構造img的變量
            label = Variable(label) 
        optimizer.zero_grad() # 消除optimizer的梯度
        out = model.forward(img) # 進行前向傳播
        loss = criterion(out, label) # 計算損失值
        loss.backward() # 進行損失值的后向傳播
        optimizer.step() # 進行優化器的優化
        train_loss += loss.data # 
    for data in test_loader:
        img, label = data
        img = img.view(img.size(0), -1)
        if torch.cuda.is_available():
            img = Variable(img, volatile=True).cuda()
            label = Variable(label, volatile=True).cuda()
        else:
            img = Variable(img, volatile=True)
            label = Variable(label, volatile=True)
        out = model.forward(img)
        loss = criterion(out, label)
        test_loss += loss.data
        top_p, top_class = out.topk(1, dim=1) # 獲得輸出的每一個樣本的最大損失
        equals = top_class == label.view(*top_class.shape) # 判斷兩組樣本的標簽是否相等
        accuracy = torch.mean(equals.type(torch.FloatTensor)) # 計算准確率 
        eval_acc += accuracy
    print('train_loss{:.6f}, test_loss{:.6f}, Acc:{:.6f}'.format(train_loss / len(train_loader), test_loss / len(test_loader), eval_acc / len(test_loader)))