Pytorch--MNIST练习--clean版

from __future__ import print_function
import sys, os, argparse
import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.autograd import Variable
import matplotlib.pyplot as plt

print(sys.path[0])

cuda = torch.cuda.is_available()
#cuda = False

parse = argparse.ArgumentParser(description='Pytorch MNIST Example')
parse.add_argument('--batchSize', type=int, default=64, metavar='input batch size')
parse.add_argument('--testBatchSize', type=int, default=1000, metavar='input batch size for testing')
parse.add_argument('--trainSize', type=int, default=10000, metavar='input dataset size(max=60000).Default=1000')
parse.add_argument('--nEpochs', type=int, default=2, metavar='number of epochs to train')
parse.add_argument('--lr', type=float, default=0.01, metavar='Learning rate.Deafault=0.01')
parse.add_argument('--momentum', type=float, default=0.5, metavar='Default=0.5',)
parse.add_argument('--seed', type=int, default=123, metavar='Romdom Seed to use.Default=123')

opt = parse.parse_args()

torch.manual_seed(opt.seed)
if cuda:
    torch.cuda.manual_seed(opt.seed)

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        root='./data',
        train=True,
        transform=torchvision.transforms.ToTensor(),
        download=True),
    batch_size=opt.batchSize,
    shuffle=True
)

test_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        root='./data',
        train=False,
        transform=torchvision.transforms.ToTensor(),
        download=True),
    batch_size=opt.batchSize,
    shuffle=True
)

print('===>Loading data')
with open('./data/processed/training.pt', 'rb') as f:
    training_set = torch.load(f)

with open('./data/processed/test.pt', 'rb') as f:
    test_set = torch.load(f)
print('<===Done')

# reshape image to 60000*1*28*28
training_data = training_set[0].view(-1, 1, 28, 28)
training_data = training_data[:opt.trainSize]
training_labels = training_set[1]
test_data = test_set[0].view(-1, 1, 28, 28)
test_labels = test_set[1]

del training_set
del test_set

# model
print('===>Building model')

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, 5)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(10, 20, 5)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.fc1   = nn.Linear(320, 50)
        self.fc2   = nn.Linear(50, 10)
        self.relu    = nn.ReLU()
        self.softmax = nn.LogSoftmax()

    def forward(self, x):
        x = self.relu(self.pool1(self.conv1(x)))
        x = self.relu(self.pool2(self.conv2(x)))
        x = x.view(-1, 320)
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        return  self.softmax(x)


model = Net()
if cuda:
    model.cuda()
    training_data   = training_data.cuda()
    training_labels = training_labels.cuda()
    test_data       = test_data.cuda()
    test_labels     = test_labels.cuda()
print('<===Done')

criterion = nn.NLLLoss()
optimizer = optim.SGD(model.parameters(), lr=opt.lr, momentum=opt.momentum)

def train(epoch):
    # create buffer for mini-batch
    batch_data   = torch.FloatTensor(opt.batchSize, 1, 28, 28)
    batch_target = torch.LongTensor(opt.batchSize)
    if cuda:
        batch_data, batch_target = batch_data.cuda(), batch_target.cuda()

    # create autograd Variable over these buffers
    batch_data   = Variable(batch_data)
    batch_target = Variable(batch_target)

    for i in range(0, training_data.size(0)-opt.batchSize+1, opt.batchSize):
        start, end = i, i+opt.batchSize
        # clear the gradients of all optimized Variables
        optimizer.zero_grad()
        batch_data.data[:]   = training_data[start:end]
        batch_target.data[:] = training_labels[start:end]
        output = model(batch_data)
        loss   = criterion(output, batch_target)
        loss.backward()
        # without this commond,loss is a Variable,con't be passed to string in later 'print'
        loss = loss.data[0]
        optimizer.step()
        print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.4f}'
              .format(epoch, end, opt.trainSize, float(end)/opt.trainSize*100, loss))


def test(epoch):
    # create buffer for mini-batch
    batch_data    = torch.FloatTensor(opt.testBatchSize, 1 , 28, 28)
    batch_targets = torch.LongTensor(opt.testBatchSize)

    if cuda:
        batch_data, batch_targets = batch_data.cuda(), batch_targets.cuda()

    # create autograd Varialble for these buffers
    batch_data    = Variable(batch_data, volatile=True)
    batch_targets = Variable(batch_targets, volatile=True)
    test_loss = 0
    correct = 0

    for i in range(0, test_data.size(0), opt.testBatchSize):
        batch_data[:]    = test_data[i:i+opt.testBatchSize]
        batch_targets[:] = test_labels[i:i+opt.testBatchSize]
        output = model(batch_data)
        test_loss += criterion(output, batch_targets)
        # get the index of the max log-probability
        pred = output.data.max(1)[1]
        correct += pred.long().eq(batch_targets.data.long()).cpu().sum()
        print('correct = ', correct)


    test_loss = test_loss.data[0]
    test_loss /= test_data.size(0)/opt.testBatchSize
    print('\nTest Set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, test_data.size(0),
        float(correct) / test_data.size(0) * 100))

for epoch in range(opt.nEpochs+1):
    train(epoch)
    test(epoch)

注：
1.在代码开头加上：from __future__ import *
自python2.1以后，当python新版本特性与老版本不兼容时，将默认禁止该特性。若想使用，则可以用该语句导入。例如我们平时用的比较多的：
from __future__ import print_function
在python2.7中加上这条语句后，我们就可以使用python3中print 函数了，代码的鲁棒性也变得更好。

2.使用view来reshape输入数据
x = torch.randn(64, 20, 4, 4)
x = x.view(64,320)
print(x.size())

([64, 320])

使用-1代表先满足其它维数据的size，最后剩下的size赋给该维：
x = torch.randn(64, 20, 4, 4)
x = x.view(-1,320)
print(x.size())

([64, 320])

需要注意的是，使用view reshape数据后，并没有保存在原来的x中，我们需要将reshape后的数据赋给x。