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

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

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
import torch
import torchvision
import torch.nn as nn
import torch.utils.data as Data
import matplotlib.pyplot as plt
from torch.autograd import Variable

# 超参数定义
EPOCH = 1
LR = 0.01
BATCH_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])
1
2
3
4
5
6
7
8
# 查看图像
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()

png

png

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
# 数据加载
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 prediction

cnn = 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)
)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
# 定义优化器
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.0982
Epoch:  0 | train loss: 0.0788 | accuracy:  0.9592
Epoch:  0 | train loss: 0.0587 | accuracy:  0.9626
Epoch:  0 | train loss: 0.0188 | accuracy:  0.9745
Epoch:  0 | train loss: 0.0707 | accuracy:  0.9759
Epoch:  0 | train loss: 0.0564 | accuracy:  0.9775
Epoch:  0 | train loss: 0.0489 | accuracy:  0.9779
Epoch:  0 | train loss: 0.0925 | accuracy:  0.9791
Epoch:  0 | train loss: 0.0566 | accuracy:  0.9834

参考资料

  1. https://www.youtube.com/user/MorvanZhou
如果有收获,可以请我喝杯咖啡!