├── README.md ├── data.rar └── test2.py /README.md: -------------------------------------------------------------------------------- 1 | # MINIST 2 | pytorch+MINIST实现手写数字识别 3 | -------------------------------------------------------------------------------- /data.rar: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jing-repo/MINIST/97063f4b640a248a95ea60f1b83b6f630e3bca23/data.rar -------------------------------------------------------------------------------- /test2.py: -------------------------------------------------------------------------------- 1 | # 训练+测试 2 | 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.utils.data as Data 7 | import torchvision 8 | import matplotlib.pyplot as plt 9 | import os 10 | import cv2 11 | 12 | torch.manual_seed(1) # 使用随机化种子使神经网络的初始化每次都相同 13 | 14 | # 超参数 15 | EPOCH = 1 # 训练整批数据的次数 16 | BATCH_SIZE = 50 17 | LR = 0.001 # 学习率 18 | DOWNLOAD_MNIST = True # 表示还没有下载数据集，如果数据集下载好了就写False 19 | 20 | # 下载mnist手写数据集 21 | train_data = torchvision.datasets.MNIST( 22 | root='./data/', # 保存或提取的位置 会放在当前文件夹中 23 | train=True, # true说明是用于训练的数据，false说明是用于测试的数据 24 | transform=torchvision.transforms.ToTensor(), # 转换PIL.Image or numpy.ndarray 25 | 26 | download=DOWNLOAD_MNIST, # 已经下载了就不需要下载了 27 | ) 28 | 29 | test_data = torchvision.datasets.MNIST( 30 | root='./data/', 31 | train=False # 表明是测试集 32 | ) 33 | 34 | # 批训练 50个samples， 1 channel，28x28 (50,1,28,28) 35 | # Torch中的DataLoader是用来包装数据的工具，它能帮我们有效迭代数据，这样就可以进行批训练 36 | train_loader = Data.DataLoader( 37 | dataset=train_data, 38 | batch_size=BATCH_SIZE, 39 | shuffle=True # 是否打乱数据，一般都打乱 40 | ) 41 | 42 | # 进行测试 43 | # 为节约时间，测试时只测试前2000个 44 | # 45 | test_x = torch.unsqueeze(test_data.train_data, dim=1).type(torch.FloatTensor)[:2000] / 255 46 | # torch.unsqueeze(a) 是用来对数据维度进行扩充，这样shape就从(2000,28,28)->(2000,1,28,28) 47 | # 图像的pixel本来是0到255之间，除以255对图像进行归一化使取值范围在(0,1) 48 | test_y = test_data.test_labels[:2000] 49 | 50 | 51 | # 用class类来建立CNN模型 52 | # CNN流程：卷积(Conv2d)-> 激励函数(ReLU)->池化(MaxPooling)-> 53 | # 卷积(Conv2d)-> 激励函数(ReLU)->池化(MaxPooling)-> 54 | # 展平多维的卷积成的特征图->接入全连接层(Linear)->输出 55 | 56 | class CNN(nn.Module): # 我们建立的CNN继承nn.Module这个模块 57 | def __init__(self): 58 | super(CNN, self).__init__() 59 | # 建立第一个卷积(Conv2d)-> 激励函数(ReLU)->池化(MaxPooling) 60 | self.conv1 = nn.Sequential( 61 | # 第一个卷积con2d 62 | nn.Conv2d( # 输入图像大小(1,28,28) 63 | in_channels=1, # 输入图片的高度，因为minist数据集是灰度图像只有一个通道 64 | out_channels=16, # n_filters 卷积核的高度 65 | kernel_size=5, # filter size 卷积核的大小 也就是长x宽=5x5 66 | stride=1, # 步长 67 | padding=2, # 想要con2d输出的图片长宽不变，就进行补零操作 padding = (kernel_size-1)/2 68 | ), # 输出图像大小(16,28,28) 69 | # 激活函数 70 | nn.ReLU(), 71 | # 池化，下采样 72 | nn.MaxPool2d(kernel_size=2), # 在2x2空间下采样 73 | # 输出图像大小(16,14,14) 74 | ) 75 | # 建立第二个卷积(Conv2d)-> 激励函数(ReLU)->池化(MaxPooling) 76 | self.conv2 = nn.Sequential( 77 | # 输入图像大小(16,14,14) 78 | nn.Conv2d( # 也可以直接简化写成nn.Conv2d(16,32,5,1,2) 79 | in_channels=16, 80 | out_channels=32, 81 | kernel_size=5, 82 | stride=1, 83 | padding=2 84 | ), 85 | # 输出图像大小 (32,14,14) 86 | nn.ReLU(), 87 | nn.MaxPool2d(2), 88 | # 输出图像大小(32,7,7) 89 | ) 90 | # 建立全卷积连接层 91 | self.out = nn.Linear(32 * 7 * 7, 10) # 输出是10个类 92 | 93 | # 下面定义x的传播路线 94 | def forward(self, x): 95 | x = self.conv1(x) # x先通过conv1 96 | x = self.conv2(x) # 再通过conv2 97 | # 把每一个批次的每一个输入都拉成一个维度，即(batch_size,32*7*7) 98 | # 因为pytorch里特征的形式是[bs,channel,h,w]，所以x.size(0)就是batchsize 99 | x = x.view(x.size(0), -1) # view就是把x弄成batchsize行个tensor 100 | output = self.out(x) 101 | return output 102 | 103 | 104 | cnn = CNN() 105 | print(cnn) 106 | 107 | # 训练 108 | # 把x和y 都放入Variable中，然后放入cnn中计算output，最后再计算误差 109 | 110 | # 优化器选择Adam 111 | optimizer = torch.optim.Adam(cnn.parameters(), lr=LR) 112 | # 损失函数 113 | loss_func = nn.CrossEntropyLoss() # 目标标签是one-hotted 114 | 115 | # 开始训练 116 | # for epoch in range(EPOCH): 117 | # for step, (b_x, b_y) in enumerate(train_loader): # 分配batch data 118 | # output = cnn(b_x) # 先将数据放到cnn中计算output 119 | # loss = loss_func(output, b_y) # 输出和真实标签的loss，二者位置不可颠倒 120 | # optimizer.zero_grad() # 清除之前学到的梯度的参数 121 | # loss.backward() # 反向传播，计算梯度 122 | # optimizer.step() # 应用梯度 123 | # 124 | # if step % 50 == 0: 125 | # test_output = cnn(test_x) 126 | # pred_y = torch.max(test_output, 1)[1].data.numpy() 127 | # accuracy = float((pred_y == test_y.data.numpy()).astype(int).sum()) / float(test_y.size(0)) 128 | # print('Epoch: ', epoch, '| train loss: %.4f' % loss.data.numpy(), '| test accuracy: %.2f' % accuracy) 129 | # 130 | # torch.save(cnn.state_dict(), 'cnn2.pkl')#保存模型 131 | 132 | # 加载模型，调用时需将前面训练及保存模型的代码注释掉，否则会再训练一遍 133 | cnn.load_state_dict(torch.load('cnn2.pkl')) 134 | cnn.eval() 135 | # print 10 predictions from test data 136 | inputs = test_x[:32] # 测试32个数据 137 | test_output = cnn(inputs) 138 | pred_y = torch.max(test_output, 1)[1].data.numpy() 139 | print(pred_y, 'prediction number') # 打印识别后的数字 140 | # print(test_y[:10].numpy(), 'real number') 141 | 142 | img = torchvision.utils.make_grid(inputs) 143 | img = img.numpy().transpose(1, 2, 0) 144 | 145 | # 下面三行为改变图片的亮度 146 | # std = [0.5, 0.5, 0.5] 147 | # mean = [0.5, 0.5, 0.5] 148 | # img = img * std + mean 149 | cv2.imshow('win', img) # opencv显示需要识别的数据图片 150 | key_pressed = cv2.waitKey(0) 151 | --------------------------------------------------------------------------------