【邱希鹏】nndl-chap5-数字识别(pytorch)-CSDN博客

本文链接：https://blog.csdn.net/Xiao_CangTian/article/details/107873977

本文介绍了使用PyTorch实现数字识别的步骤，包括导入所需包、加载数据、构建CNN模型，以及详细展开的训练与测试过程。在训练和测试阶段，展示了数据集的尺寸和部分预测结果。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

1. 导入包

import os
import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.utils.data as Data
import torchvision
import torch.nn.functional as F
import numpy as np
from torchvision import datasets, transforms
learning_rate = 1e-4
keep_prob_rate = 0.7 #
max_epoch = 3
BATCH_SIZE = 50

DOWNLOAD_MNIST = False
if not(os.path.exists('MNIST')) or not os.listdir('MNIST'):
    # not mnist dir or mnist is empyt dir
    DOWNLOAD_MNIST = True

2. 导入数据

train_data = torchvision.datasets.MNIST(root='./', train = True, download=DOWNLOAD_MNIST,
                                        transform = torchvision.transforms.Compose([
                                            transforms.ToTensor(),
                                        ]))
train_loader = Data.DataLoader(dataset = train_data ,batch_size= BATCH_SIZE ,shuffle= True)

test_data = torchvision.datasets.MNIST(root = './', train = False,
                                      transform = transforms.Compose([
                                          transforms.ToTensor(),
                                      ]))
test_loader = Data.DataLoader(dataset = test_data, batch_size = BATCH_SIZE, shuffle = True)

test_x = Variable(torch.unsqueeze(test_data.test_data,dim  = 1),volatile = True).type(torch.FloatTensor)[:500]/255.
test_y = test_data.test_labels[:500].numpy()
print(test_x.shape)
print(test_y.shape)

# torch.Size([500, 1, 28, 28])
# (500,)

3. CNN模型

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d( # ???
                # patch 7 * 7 ; 1  in channels ; 32 out channels ; ; stride is 1
                # padding style is same(that means the convolution opration's input and output have the same size)
                in_channels = 1     ,  
                out_channels = 32   ,
                kernel_size = 7     ,
                stride = 1          ,
                padding = 0    ,
            ),
            nn.ReLU(),        # activation function
            nn.MaxPool2d(2),  # pooling operation
        )
        self.conv2 = nn.Sequential( # ???
            # line 1 : convolution function, patch 5*5 , 32 in channels ;64 out channels; padding style is same; stride is 1
            # line 2 : choosing your activation funciont
            # line 3 : pooling operation function.
            nn.Conv2d(
                in_channels = 32,
                out_channels = 64,
                kernel_size = 5,
                stride = 1,
                padding = 0,
            ),
            nn.ReLU(),
            nn.MaxPool2d(1),
        )
        self.out1 = nn.Linear(7*7*64 , 1024 , bias= True)   # full connection layer one

        self.dropout = nn.Dropout(keep_prob_rate)
        self.out2 = nn.Linear(1024, 10, bias=True)


    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(-1, 64*7*7)  # flatten the output of coonv2 to (