直接代码：

import torch
import torch.nn as nn
import torch.utils.data as Data
import torchvision
import matplotlib.pyplot as plt


# Hyper parameters
torch.manual_seed(1)
EPOCH = 1  #训练一批次
BATCH_SICE = 50  # 一次50个
LR = 0.001  # 每次的学习率
DOWNLOAD_MNSIT = False  # 需要下载为True， 下载好了可以设置为False

# download data
train_data = torchvision.datasets.MNIST(
    root='./mnist/',
    train=True,
    transform=torchvision.transforms.ToTensor(),
    download=DOWNLOAD_MNSIT
)

test_data = torchvision.datasets.MNIST(
    root='./mnist/',
    train=False
)

# train_loader
train_loader = Data.DataLoader(
    dataset=train_data,
    batch_size=BATCH_SICE,
    shuffle=True
)

# ready_for_test 前20000个
test_x = torch.unsqueeze(test_data.test_data, dim=1).type(torch.Tensor)[:2000]/255.
test_y = test_data.test_labels[:2000]
# print(test_x.size())
# print(test_y.size())
# print(test_x[1])
# print(test_y[1])

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1  = nn.Sequential(     # (1, 28, 28)
            nn.Conv2d(
                in_channels=1,
                out_channels=16,
                kernel_size=5,
                stride=1,
                padding=2
            ),             # output shape : N = (W - F + 2P)/S + 1
            # W :图片大小 W*W 28, F:filter大小 F*F，5， S：步长, 1  P:padding 2
            # N = (28 - 5 + 2*2)/1 + 1 = 28 因为有16个filter
            # (16, 28, 28)
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
            # Maxpooling 2*2 向下采样取最大（想想如果4*4图片，采用2*2maxpooling，就变成了2*2图片）
            # 所以这里图片大小为 (16, 14, 14)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 5, 1, 2),  # (32, 14, 14)
            nn.ReLU(),
            nn.MaxPool2d(2)  # (32, 7, 7)
        )
        self.out = nn.Linear(32 * 7 * 7, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)  # 展平多维的卷积图成 (batch_size, 32 * 7 * 7)
        output = self.out(x)
        return output

cnn = CNN()
print(cnn)



# training
optimizer = torch.optim.Adam(cnn.parameters(), lr=LR)
loss_func = nn.CrossEntropyLoss()


# training and testing
for epoch in range(EPOCH):
    for step, (b_x, b_y) in enumerate(train_loader):   # 分配 batch data, normalize x when iterate train_loader
        output = cnn(b_x)               # cnn output
        loss = loss_func(output, b_y)   # cross entropy loss
        optimizer.zero_grad()           # clear gradients for this training step
        loss.backward()                 # backpropagation, compute gradients
        optimizer.step()                # apply gradients

        if step % 50 == 0:
            print('train loss = ', loss.data.numpy())

结果

CNN(
  (conv1): Sequential(
    (0): Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (conv2): Sequential(
    (0): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (out): Linear(in_features=1568, out_features=10, bias=True)
)
train loss =  2.3104544
train loss =  0.61840093
train loss =  0.12703253
train loss =  0.23725168
train loss =  0.4044273
train loss =  0.08451731
train loss =  0.19528298
train loss =  0.109065436
train loss =  0.123532385
train loss =  0.06730688
train loss =  0.2240683
train loss =  0.2114728
train loss =  0.024014007
train loss =  0.08469809
train loss =  0.21586336
train loss =  0.10181876
train loss =  0.043114547
train loss =  0.09106462
train loss =  0.055737924
train loss =  0.10089029
train loss =  0.032855053
train loss =  0.021929108
train loss =  0.025414439
train loss =  0.117736794

Process finished with exit code 0