1. import torch
    2. import torch.nn as nn
    3. import torch.utils.data as Data
    4. import torchvision      
    5. import matplotlib.pyplot as plt
    7. torch.manual_seed(1)    # reproducible
    9. # Hyper Parameters
    10. EPOCH = 5           # 训练整批数据多少次, 为了节约时间, 我们只训练一次
    11. BATCH_SIZE = 50
    12. LR = 0.001          # 学习率
    13. DOWNLOAD_MNIST = True  # 如果你已经下载好了mnist数据就写上 False
    16. # Mnist 手写数字
    17. train_data = torchvision.datasets.MNIST(
    18.     root='mnist',    # 保存或者提取位置
    19.     train=True,  # this is training data
    20.     transform=torchvision.transforms.ToTensor(),   
    21.     download=DOWNLOAD_MNIST,         
    22. )
    24. test_data = torchvision.datasets.MNIST(root='mnist', train=False)
    26. # 批训练 50samples, 1 channel, 28x28 (50, 1, 28, 28)
    27. train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)
    29. # 为了节约时间, 我们测试时只测试前2000个
    30. test_x = torch.unsqueeze(test_data.data, dim=1).type(torch.FloatTensor)[:2000]/255.   # shape from (2000, 28, 28) to (2000, 1, 28, 28), value in range(0,1)
    31. test_y = test_data.targets[:2000]
    33. class CNN(nn.Module):
    34.     def __init__(self):
    35.         super(CNN, self).__init__()
    36.         self.conv1 = nn.Sequential(  # input shape (1, 28, 28)
    37.             nn.Conv2d(
    38.                 in_channels=1,      # input height
    39.                 out_channels=16,    # n_filters
    40.                 kernel_size=5,      # filter size
    41.                 stride=1,           # filter movement/step
    42.                 padding=2,      # 如果想要 con2d 出来的图片长宽没有变化, padding=(kernel_size-1)/2 当 stride=1
    43.             ),      # output shape (16, 28, 28)
    44.             nn.ReLU(),    # activation
    45.             nn.MaxPool2d(kernel_size=2),    # 在 2x2 空间里向下采样, output shape (16, 14, 14)
    46.         )
    47.         self.conv2 = nn.Sequential(  # input shape (16, 14, 14)
    48.             nn.Conv2d(16, 32, 5, 1, 2),  # output shape (32, 14, 14)
    49.             nn.ReLU(),  # activation
    50.             nn.MaxPool2d(2),  # output shape (32, 7, 7)
    51.         )
    52.         self.out = nn.Linear(32 * 7 * 7, 10)   # fully connected layer, output 10 classes
    54.     def forward(self, x):
    55.         x = self.conv1(x)
    56.         x = self.conv2(x)
    57.         x = x.view(x.size(0), -1)   # 展平多维的卷积图成 (batch_size, 32 * 7 * 7)
    58.         output = self.out(x)
    59.         return output
    61. cnn = CNN()
    62. print(cnn)  # net architecture
    63. cnn=cnn.cuda()
    64. optimizer = torch.optim.Adam(cnn.parameters(), lr=LR)   # optimize all cnn parameters
    65. loss_func = nn.CrossEntropyLoss()   # the target label is not one-hotted
    67. # training and testing
    68. for epoch in range(EPOCH):
    69.     to=0
    70.     ac=0
    71.     for step, (b_x, b_y) in enumerate(train_loader):   # 分配 batch data, normalize x when iterate train_loader
    72.         b_x=b_x.cuda()
    73.         b_y=b_y.cuda()
    74.         output = cnn(b_x)               # cnn output
    76.         loss = loss_func(output, b_y)   # cross entropy loss
    77.         optimizer.zero_grad()           # clear gradients for this training step
    78.         loss.backward()                 # backpropagation, compute gradients
    79.         optimizer.step()                # apply gradients
    80.         pred_y = torch.max(output, 1)[1].data.cpu().numpy().squeeze()
    81.         acc_y=b_y.data.cpu().numpy()
    82.         for i in range(len(pred_y)):
    83.             if acc_y[i]==pred_y[i]:
    84.                 ac+=1
    85.             to+=1
    86.     print("Epoch:",epoch+1,"Acc",ac/to)
    88. test=test_x[:10].cuda()
    89. test_output = cnn(test)
    90. pred_y = torch.max(test_output, 1)[1].data.cpu().numpy().squeeze()
    91. print(pred_y, 'prediction number')
    92. print(test_y[:10].cpu().numpy(), 'real number')

    image.png