pytorch笔记9--CNN & GPU加速

数字识别 

注：用GPU加速时： 1. 将数据放的GPU上  2. 将网络放到GPU上  3. 将相应的操作放到GPU上

# 数字识别
#  用GPU加速时： 1.将数据放的GPU上 2.将网络放到GPU上 3. 将相应的操作放到GPU上

import torch

import torch.nn as nn

import torch.utils.data as Data

import torchvision

import matplotlib.pyplot as plt

torch.manual_seed(1)

# Hyper Parameters

EPOCH = 1

BATCH_SIZE = 50

LR = 0.001

DOWNLOAD_MNIST = False

# Mnist数据集下载

train_data = torchvision.datasets.MNIST(

    root='./mnist/',  # 保存或提取的位置

    train=True,

    transform=torchvision.transforms.ToTensor(),  # 将下载的数据集PIL.Image or numpy.ndarray

    # 转换为torch.FloatTensor(c*h*w),训练的时候normalize成[0.0,1.0]区间

    download=DOWNLOAD_MNIST,  # 下不下载

)

# plot one example

# print(train_data.train_data.size())

# print(train_data.test_data.size())

# plt.imshow(train_data.train_data[0].numpy(),cmap='gray')

# plt.title('%i'%train_data.train_labels[0])

# plt.show()

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

# 批训练 50samples,1 channel,28*28 (50,1,28,28)

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

# 取前2000个数据进行测试 shape from (2000,28,28) to (2000,1,28,28), 值在[0.0,1.0]


'''
GPU   对数据的处理
'''
test_x = torch.unsqueeze(test_data.data, dim=1).type(torch.FloatTensor)[:2000].cuda() / 255
test_y = test_data.targets[:2000]


# CNN

class CNN(nn.Module):

    def __init__(self):
        super(CNN, self).__init__()

        self.conv1 = nn.Sequential(

            nn.Conv2d(  # 输入：(1,28,28)

                in_channels=1,

                out_channels=16,

                kernel_size=5,

                stride=1,

                padding=2,  # 如果stride=1,卷积后要使得图片的长宽不变，padding=(kernel_size-1)/2

            ),  # 输出：(16,28,28)

            nn.ReLU(),

            nn.MaxPool2d(kernel_size=2),  # (16,14,14)

        )

        self.conv2 = nn.Sequential(

            nn.Conv2d(16, 32, 5, 1, 2),

            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)

        output = self.out(x)

        return output


cnn = CNN()
'''
GPU  对网络的处理
'''
cnn.cuda()

# 训练

optimizer = torch.optim.Adam(cnn.parameters(), lr=LR)

loss_func = nn.CrossEntropyLoss()

for epoch in range(EPOCH):

    for step, (b_x, b_y) in enumerate(train_loader):
        '''
        GPU  对数据的处理
        '''
        b_x=b_x.cuda()
        b_y=b_y.cuda()

        output = cnn(b_x)

        loss = loss_func(output, b_y)

        optimizer.zero_grad()

        loss.backward()

        optimizer.step()

        if step % 50 == 0:
            test_out = cnn(test_x)

            '''
            GPU  将操作放到GPU上加速
            '''
            pred_y = torch.max(test_out, 1)[1].cuda().data.squeeze()  # torch.max(input,dim)返回每一行中的最大值的标签

            '''
            CPU  上步得到的pred_y是cuda的形式,  但是can't convert CUDA tensor to numpy
            '''
            pred_y=pred_y.cpu()
            accuracy = (pred_y == test_y).numpy().sum() / test_y.size(0)

            print('step: {} | train loss: {} | test accuracy: {} '.format(step, loss.data, accuracy))

# 打印前十个数据的预测结果和真实结果，看一下效果

test_output = cnn(test_x[:10])

'''
GPU  将操作放到 GPU
'''
pred_label = torch.max(test_output, 1)[1].cuda()

'''
CPU  ,can't convert CUDA tensor to numpy
'''
pred_label=pred_label.cpu().data.numpy()

print('prediction:  ', pred_label)

print('real:        ', test_y[:10].numpy())