pytorch实现softmax多分类（MNIST）

#导入各种库
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# Training settings
batch_size = 64

#数据集的处理

#图片变换:转换成Tensor,标准化
transform = transforms.Compose([ transforms.ToTensor(),
                                transforms.Normalize((0.1307, ),(0.3081, ))])
#创建训练数据集
train_dataset = datasets.MNIST(root='./mnist_data/',
                               train=True, download=True,
                               transform=transform)
#导入训练数据集
train_loader = DataLoader(train_dataset,
                          shuffle=True,
                          batch_size=batch_size)
#创建测试数据集
test_dataset = datasets.MNIST(root='./mnist_data/',
                              train=False,
                              download=True,
                              transform=transform)
#导入测试数据集
test_loader = DataLoader(test_dataset,
                          shuffle=False,
                          batch_size=batch_size)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.l1 = torch.nn.Linear(784, 512)
        self.l2 = torch.nn.Linear(512, 256)
        self.l3 = torch.nn.Linear(256, 128)
        self.l4 = torch.nn.Linear(128, 64)
        self.l5 = torch.nn.Linear(64, 10)

    def forward(self, x):
        # Flatten the data (n, 1, 28, 28) --> (n, 784)
        x = x.view(-1, 784)
        x = F.relu(self.l1(x))
        x = F.relu(self.l2(x))
        x = F.relu(self.l3(x))
        x = F.relu(self.l4(x))
        return F.log_softmax(self.l5(x), dim=1)

model = Net()

criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)

def train(epoch):
    # 每次输入barch_idx个数据
    for batch_idx, data in enumerate(train_loader):
        inputs, target = data
        optimizer.zero_grad()
        output = model(inputs)
        # loss
        loss = criterion(output, target)
        loss.backward()
        # update
        optimizer.step()
        if batch_idx % 200 == 0:   #len(data)=64，理解为batch-size,len(train_loader.dataset)=60000总样本数，len(train_loader)是有多少个loader，理解为共有多少个iterations
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(inputs), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))

def test():
    test_loss = 0
    correct = 0
    # 测试集
    # 不计算梯度，节省内存
    with torch.no_grad():
        for data in test_loader:
            images, labels = data
            output = model(images)
            # sum up batch loss
            test_loss += criterion(output, labels).item()
            # get the index of the max
            pred = output.data.max(1, keepdim=True)[1]
            correct += pred.eq(labels.data.view_as(pred)).cpu().sum()

    test_loss /= len(test_loader.dataset)
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

for epoch in range(1,10):
    train(epoch)
    test()

log如下：
/usr/local/bin/python3.6 /Users/function/Documents/code/test2.py
Train Epoch: 1 [0/60000 (0%)] Loss: 2.309441
Train Epoch: 1 [12800/60000 (21%)] Loss: 2.183163
Train Epoch: 1 [25600/60000 (43%)] Loss: 0.835843
Train Epoch: 1 [38400/60000 (64%)] Loss: 0.439794
Train Epoch: 1 [51200/60000 (85%)] Loss: 0.261408

Test set: Average loss: 0.0057, Accuracy: 8866/10000 (89%)

Train Epoch: 2 [0/60000 (0%)] Loss: 0.508892
Train Epoch: 2 [12800/60000 (21%)] Loss: 0.333321
Train Epoch: 2 [25600/60000 (43%)] Loss: 0.200243
Train Epoch: 2 [38400/60000 (64%)] Loss: 0.277475
Train Epoch: 2 [51200/60000 (85%)] Loss: 0.187367

Test set: Average loss: 0.0030, Accuracy: 9425/10000 (94%)

Train Epoch: 3 [0/60000 (0%)] Loss: 0.145257
Train Epoch: 3 [12800/60000 (21%)] Loss: 0.043807
Train Epoch: 3 [25600/60000 (43%)] Loss: 0.183113
Train Epoch: 3 [38400/60000 (64%)] Loss: 0.186657
Train Epoch: 3 [51200/60000 (85%)] Loss: 0.101876

Test set: Average loss: 0.0022, Accuracy: 9578/10000 (96%)

Train Epoch: 4 [0/60000 (0%)] Loss: 0.184665
Train Epoch: 4 [12800/60000 (21%)] Loss: 0.340358
Train Epoch: 4 [25600/60000 (43%)] Loss: 0.151537
Train Epoch: 4 [38400/60000 (64%)] Loss: 0.025888
Train Epoch: 4 [51200/60000 (85%)] Loss: 0.175476

Test set: Average loss: 0.0020, Accuracy: 9610/10000 (96%)

Train Epoch: 5 [0/60000 (0%)] Loss: 0.112639
Train Epoch: 5 [12800/60000 (21%)] Loss: 0.130508
Train Epoch: 5 [25600/60000 (43%)] Loss: 0.107424
Train Epoch: 5 [38400/60000 (64%)] Loss: 0.023164
Train Epoch: 5 [51200/60000 (85%)] Loss: 0.075948

Test set: Average loss: 0.0017, Accuracy: 9673/10000 (97%)

Train Epoch: 6 [0/60000 (0%)] Loss: 0.037222
Train Epoch: 6 [12800/60000 (21%)] Loss: 0.033421
Train Epoch: 6 [25600/60000 (43%)] Loss: 0.019519
Train Epoch: 6 [38400/60000 (64%)] Loss: 0.229080
Train Epoch: 6 [51200/60000 (85%)] Loss: 0.012769

Test set: Average loss: 0.0015, Accuracy: 9698/10000 (97%)

Train Epoch: 7 [0/60000 (0%)] Loss: 0.024760
Train Epoch: 7 [12800/60000 (21%)] Loss: 0.043735
Train Epoch: 7 [25600/60000 (43%)] Loss: 0.025668
Train Epoch: 7 [38400/60000 (64%)] Loss: 0.030804
Train Epoch: 7 [51200/60000 (85%)] Loss: 0.048212

Test set: Average loss: 0.0014, Accuracy: 9743/10000 (97%)

Train Epoch: 8 [0/60000 (0%)] Loss: 0.061919
Train Epoch: 8 [12800/60000 (21%)] Loss: 0.047970
Train Epoch: 8 [25600/60000 (43%)] Loss: 0.088634
Train Epoch: 8 [38400/60000 (64%)] Loss: 0.069933
Train Epoch: 8 [51200/60000 (85%)] Loss: 0.010266

Test set: Average loss: 0.0013, Accuracy: 9739/10000 (97%)

Train Epoch: 9 [0/60000 (0%)] Loss: 0.024968
Train Epoch: 9 [12800/60000 (21%)] Loss: 0.024904
Train Epoch: 9 [25600/60000 (43%)] Loss: 0.071306
Train Epoch: 9 [38400/60000 (64%)] Loss: 0.003656
Train Epoch: 9 [51200/60000 (85%)] Loss: 0.023971

Test set: Average loss: 0.0012, Accuracy: 9768/10000 (98%)