PyTorch基本用法(十)——卷积神经网络

文章作者：Tyan
博客：noahsnail.com  |  CSDN  |  简书

本文主要是关于PyTorch的一些用法。

import torchimport torchvisionimport torch.nn as nnimport torch.utils.data as Dataimport matplotlib.pyplot as pltfrom torch.autograd import Variable# 超参数定义EPOCH = 1LR = 0.01BATCH_SIZE = 64# 下载MNIST数据集train_data = torchvision.datasets.MNIST(    root = './mnist/',    # 是否是训练数据    train = True,    # 数据变换(0, 255) -> (0, 1)    transform = torchvision.transforms.ToTensor(),    # 是否下载MNIST数据    download = True)test_data = torchvision.datasets.MNIST(    root = './mnist/',    # 是否是训练数据    train = False,    # 数据变换(0, 255) -> (0, 1)    transform = torchvision.transforms.ToTensor(),    # 是否下载MNIST数据    download = True)print train_data.train_data.size()print train_data.train_labels.size()print test_data.test_data.size()print test_data.test_labels.size()

torch.Size([60000, 28, 28])torch.Size([60000])torch.Size([10000, 28, 28])torch.Size([10000])

# 查看图像plt.imshow(train_data.train_data[0].numpy(), cmap = 'gray')plt.title('%i' % train_data.train_labels[0])plt.show()plt.imshow(test_data.test_data[0].numpy(), cmap = 'gray')plt.title('%i' % test_data.test_labels[0])plt.show()

# 数据加载train_loader = Data.DataLoader(dataset = train_data, batch_size = BATCH_SIZE, shuffle = True, num_workers = 2)test_loader = Data.DataLoader(dataset = test_data, batch_size = BATCH_SIZE, shuffle = False, num_workers = 1)# 定义卷积神经网络class CNN(nn.Module):    def __init__(self):        super(CNN, self).__init__()        self.conv1 = nn.Sequential(            nn.Conv2d(                in_channels = 1,                out_channels = 16,                kernel_size = 5,                stride = 1,                padding = 2            ),            nn.ReLU(),            nn.MaxPool2d(kernel_size = 2)        )        # conv1输出为(16, 14, 14)        self.conv2 = nn.Sequential(            nn.Conv2d(16, 32, 5, 1, 2),            nn.ReLU(),            nn.MaxPool2d(2)        )        # conv2输出为(32, 7, 7)        self.output = 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)        prediction = self.output(x)        return predictioncnn = CNN()print cnn

CNN (  (conv1): Sequential (    (0): Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))    (1): ReLU ()    (2): MaxPool2d (size=(2, 2), stride=(2, 2), dilation=(1, 1))  )  (conv2): Sequential (    (0): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))    (1): ReLU ()    (2): MaxPool2d (size=(2, 2), stride=(2, 2), dilation=(1, 1))  )  (output): Linear (1568 -> 10))

# 定义优化器optimizer = torch.optim.Adam(cnn.parameters(), lr = LR, betas= (0.9, 0.999))# 定义损失函数loss_func = nn.CrossEntropyLoss()# 训练for epoch in xrange(EPOCH):    for step, (x, y) in enumerate(train_loader):        x_var = Variable(x)        y_var = Variable(y)        prediction = cnn(x_var)        loss = loss_func(prediction, y_var)        optimizer.zero_grad()        loss.backward()        optimizer.step()        if step % 100 == 0:            correct = 0.0            for step_test, (test_x, test_y) in enumerate(test_loader):                test_x = Variable(test_x)                test_output = cnn(test_x)                pred_y = torch.max(test_output, 1)[1].data.squeeze()                correct += sum(pred_y == test_y)            accuracy = correct / test_data.test_data.size(0)            print 'Epoch: ', epoch, '| train loss: %.4f' % loss.data[0], '| accuracy: ', accuracy

Epoch:  0 | train loss: 2.2787 | accuracy:  0.0982Epoch:  0 | train loss: 0.0788 | accuracy:  0.9592Epoch:  0 | train loss: 0.0587 | accuracy:  0.9626Epoch:  0 | train loss: 0.0188 | accuracy:  0.9745Epoch:  0 | train loss: 0.0707 | accuracy:  0.9759Epoch:  0 | train loss: 0.0564 | accuracy:  0.9775Epoch:  0 | train loss: 0.0489 | accuracy:  0.9779Epoch:  0 | train loss: 0.0925 | accuracy:  0.9791Epoch:  0 | train loss: 0.0566 | accuracy:  0.9834