PyTorch MNIST实验

#!/root/anaconda3/bin/python
import os

import torch
import torchvision
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision import transforms
from torch.autograd import Variable
from torch import optim
import torch.nn as nn
import torch.nn.functional as F

learning_rate = 1e-3
batch_size = 100
epoches = 5

trans_img = transforms.Compose([
        transforms.ToTensor()
    ])

trainset = MNIST('./data', train=True, transform=trans_img)
testset = MNIST('./data', train=False, transform=trans_img)

trainloader = DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=4)
testloader = DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=4)

class Lenet(nn.Module):
    def __init__(self):
        super(Lenet, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 6, 3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(6, 16, 5, stride=1, padding=0),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )

# 这个地方要根据输入shape[28, 28]自己计算全连接层的输入维度
# 计算结果是400(16*5*5),MXNet的Gluon可以实现自动计算
        self.fc = nn.Sequential(
            nn.Linear(400, 120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, 10)
        )

    def forward(self, x):
        out = self.conv(x)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out

lenet = Lenet()

criterian = nn.CrossEntropyLoss(size_average=False)
optimizer = optim.SGD(lenet.parameters(), lr=learning_rate)

# train
for i in range(epoches):
    running_loss = 0.
    running_acc = 0.
    for (img, label) in trainloader:
        img = Variable(img)
        label = Variable(label)

        optimizer.zero_grad()
        output = lenet(img)
        loss = criterian(output, label)

        # backward
        loss.backward()
        optimizer.step()

        running_loss += loss.data[0]
        _, predict = torch.max(output, 1)

#       print(label)
#       os._exit(1)

        correct_num = (predict == label).sum()
        running_acc += correct_num.data[0]

    running_loss /= len(trainset)
    running_acc /= len(trainset)

    print("[%d/%d] Loss: %.5f, Acc: %.2f" %(i+1, epoches, running_loss, 100*running_acc))

# 针对包含dropout或批量归一化层的模型
lenet.eval()

testloss = 0.
testacc = 0.
for (img, label) in testloader:
    img = Variable(img)
    label = Variable(label)

    output = lenet(img)
    loss = criterian(output, label)
    testloss += loss.data[0]
    _, predict = torch.max(output, 1)
    num_correct = (predict == label).sum()
    testacc += num_correct.data[0]

testloss /= len(testset)
testacc /= len(testset)

print(len(testset))
print("Test: Loss: %.5f, Acc: %.2f" % (testloss, 100*testacc))
参考文章
https://zhuanlan.zhihu.com/p/26635366


发表于:2017-11-16 17:54:24

原文链接(转载请保留): http://www.multisilicon.com/blog/a24346427.html

友情链接: MICROIC
首页