pytorch 用卷积神经网络做数字识别
虽然说用卷积来做数字识别实在是大材小用,但用来学习pytorch和卷积神经网络还是不错的。(ps:还是我的电脑太low)
附上我的网络:
class CNN(nn.Module):
def __init__(self,in_dim,n_class):
super(CNN,self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(in_dim,6,kernel_size=3,stride=1,padding=1),
# input shape(1*28*28),(28+1*2-3)/1+1=28 卷积后输出(6*28*28)
# 输出图像大小计算公式:(n*n像素的图)(n+2p-k)/s+1
nn.ReLU(True), # 激活函数
nn.MaxPool2d(2,2), # 28/2=14 池化后(6*14*14)
nn.Conv2d(6,16,5,stride=1,padding=0), # (14-5)/1+1=10 卷积后(16*10*10)
nn.ReLU(True),
nn.MaxPool2d(2,2) #池化后(16*5*5)=400,the input of full connection
)
self.fc = nn.Sequential( #full connection layers.
nn.Linear(400,120),
nn.Linear(120,84),
nn.Linear(84,n_class)
)
def forward(self, x):
out = self.conv(x) #out shape(batch,16,5,5)
out = out.view(out.size(0), -1) #out shape(batch,400)
out = self.fc(out) #out shape(batch,10)
return out
net = CNN(1, 10)
print(net)
方便管理我将代码分为train.py test.py tester.py 你问我tester.py是啥。我随便写个数字来做一下检验。如下:
图片处理过程:
input_image = 'D:/4.png'
im = Image.open(input_image).resize((28, 28)) #取图片数据
im = im.convert('L') #灰度图
im_data = np.array(im)
im_data = torch.from_numpy(im_data).float()
im_data = im_data.view(1, 1, 28, 28)
识别结果:
附上我的全部代码链接:github