pytorch学习笔记---搭建CNN识别MNIST

CNN网络搭建

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

# define hyper parameters
Batch_size = 100
Epoch = 1
Lr = 0.5

# define train data and test data
train_data = torchvision.datasets.MNIST(
    root='./mnist',
    train=True,
    download=False,
    transform=torchvision.transforms.ToTensor()
)
# print(train_data.data.shape)           # torch.Size([60000, 28, 28])
# print(train_data.targets.size())        # torch.Size([60000])
# print(train_data.data[0].size())        # torch.Size([28, 28])
# plt.imshow(train_data.data[0].numpy(), cmap='gray')     
# plt.show()
test_data = torchvision.datasets.MNIST(
    root='./mnist',
    train=False,
    # transform=torchvision.transforms.ToTensor()
)
test_x = torch.unsqueeze(test_data.data, dim=1).type(torch.FloatTensor)[:2000]
test_y = test_data.targets[:2000]
# print(test_x.shape)                         # torch.Size([2000, 1, 28, 28])
# print(test_y.shape)                         # torch.Size([2000])
train_loader = Data.DataLoader(
    dataset=train_data,
    shuffle=True,
    batch_size=Batch_size,
)

# define network structure
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.con1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.con2 = nn.Sequential(
            nn.Conv2d(16, 32, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.out = nn.Linear(32 * 7 * 7, 10)

    def forward(self, x):
        x = self.con1(x)            # (batch, 16, 14, 14)
        x = self.con2(x)            # (batch, 32, 7, 7)
        x = x.view(x.size(0), -1)
        out = self.out(x)             # (batch_size, 10)
        return out

cnn = CNN()
# print(cnn)
optimizer = torch.optim.SGD(cnn.parameters(), lr=Lr)
loss_fun = nn.CrossEntropyLoss()

for epoch in range(Epoch):
    for i, (x, y) in enumerate(train_loader):
        output = cnn(x)
        loss = loss_fun(output, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if i % 50 == 0:
            test_output = torch.max(cnn(test_x), dim=1)[1]
            loss = loss_fun(cnn(test_x), test_y).item()
            accuracy = torch.sum(torch.eq(test_output, test_y)).item() / test_y.numpy().size
            print('Epoch:', Epoch, '|loss:%.4f' % loss, '|accuracy:%.4f' % accuracy)

print('real value', test_data.targets[: 10].numpy())
print('train value', torch.max(cnn(test_x)[: 10], dim=1)[1].numpy())

最后输出结果为：

Epoch: 1 |loss:38.7337 |accuracy:0.1025
Epoch: 1 |loss:198.9613 |accuracy:0.5905
Epoch: 1 |loss:45.0082 |accuracy:0.9095
Epoch: 1 |loss:23.3548 |accuracy:0.9500
Epoch: 1 |loss:20.2242 |accuracy:0.9545
Epoch: 1 |loss:19.9034 |accuracy:0.9580
Epoch: 1 |loss:29.5114 |accuracy:0.9475
Epoch: 1 |loss:16.1782 |accuracy:0.9645
Epoch: 1 |loss:14.2296 |accuracy:0.9665
Epoch: 1 |loss:13.3096 |accuracy:0.9660
Epoch: 1 |loss:15.6233 |accuracy:0.9635
Epoch: 1 |loss:21.5307 |accuracy:0.9560
real value [7 2 1 0 4 1 4 9 5 9]
train value [7 2 1 0 4 1 4 9 5 9]

注意点

在搭建网络时，有几点需要特别注意的地方。

1、TypeError:default_collate: batch must contain tensors, numpy arrays, numbers, dicts or lists; found <class ‘PIL.Image.Image’>

1. 对于train_data 使用batch_size的方式进行训练，一定要设置batch的类型，所以规定train_data一定要加语句transform=torchvision.transforms.ToTensor()

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

否则就会出现错误：

TypeError: default_collate: batch must contain tensors, numpy arrays, numbers, dicts or lists; found <class 'PIL.Image.Image'>

2、RuntimeError: Expected 4-dimensional input for 4-dimensional weight [16, 1, 5, 5], but got 3-dimensional input of size [2000, 28, 28] instead

在test_data时不需要batch，从测试集中选取部分测试集检验训练效果。所以在设置test_data时不需要规定：transform=torchvision.transforms.ToTensor()
但是，这样的话相当于在训练CNN时少了一个维度，就会出现错误：

RuntimeError: Expected 4-dimensional input for 4-dimensional weight [16, 1, 5, 5], but got 3-dimensional input of size [2000, 28, 28] instead

所以为了不出现维度错误，必须在训练时给test_x增加一个维度，要用语句：

test_x = torch.unsqueeze(test_data.data, dim=1).type(torch.FloatTensor)[:2000]