# PyTorch：基于CNN的数字识别(MNIST)

## 背景

train_gpu.py

import torch
from torch import nn
from torch import optim
import torch.nn.functional as F
from torchvision import datasets, transforms

# 我设计的网络结构，包括三个卷积层和三个全链接层
class myModel(nn.Module):
def __init__(self):
super(myModel, self).__init__()
self.conv1 = nn.Conv2d(1, 64, kernel_size=5, stride=2, padding=1)
self.conv2 = nn.Conv2d(64, 192, kernel_size=3, stride=2, padding=1)
self.conv3 = nn.Conv2d(192, 384, kernel_size=3, stride=2, padding=1)
self.fc1   = nn.Linear( 6144,4096)
self.fc2   = nn.Linear( 4096,4096)
self.fc3   = nn.Linear( 4096,10)
self.dropout = nn.Dropout(0.25)
self.softmax = nn.LogSoftmax(dim=1)

def forward(self, x):
# forward
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = F.relu(self.conv3(x))
x = x.view(x.size(0), 6144 )
nn.Dropout()
x = F.relu(self.fc1(x))
nn.Dropout()
x = F.relu(self.fc2(x))
nn.Dropout()
x = F.relu(self.fc3(x))
x = self.softmax(x)
return x

# Define a transform to normalize the data
transform = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,)),
])

model = myModel()
model.cuda()

criterion = nn.NLLLoss()
optimizer = optim.SGD(model.parameters(), lr=0.003)

epochs = 100
train_losses, test_losses = [], []

for e in range(epochs):
running_loss = 0
# Flatten MNIST images into a 784 long vector
images = images.cuda()
labels = labels.cuda()

# TODO: Training pass

output = model(images)
loss = criterion(output, labels)
loss.backward()
optimizer.step()

running_loss += loss.item()
else:
test_loss = 0
accuracy = 0

# Turn off gradients for validation, saves memory and computations
images = images.cuda()
labels = labels.cuda()

log_ps = model(images)
test_loss += criterion(log_ps, labels)

ps = torch.exp(log_ps)
top_p, top_class = ps.topk(1, dim=1)
equals = top_class == labels.view(*top_class.shape)
accuracy += torch.mean(equals.type(torch.FloatTensor))

torch.save(model.state_dict(), str(e+1) +'.pth')

print("Epoch: {}/{}.. ".format( e+1, epochs),

import matplotlib.pyplot as plt
plt.plot(train_losses, label='Training loss')
plt.plot(test_losses, label='Validation loss')
plt.legend(frameon=False)



## 二、预测

predict.py

import torch
from torch import nn
from torch import optim
import torch.nn.functional as F
from torchvision import datasets, transforms

# define the CNN architecture
class myModel(nn.Module):
def __init__(self):
super(myModel, self).__init__()
self.conv1 = nn.Conv2d(1, 64, kernel_size=5, stride=2, padding=1)
self.conv2 = nn.Conv2d(64, 192, kernel_size=3, stride=2, padding=1)
self.conv3 = nn.Conv2d(192, 384, kernel_size=3, stride=2, padding=1)
self.fc1   = nn.Linear( 6144,4096)
self.fc2   = nn.Linear( 4096,4096)
self.fc3   = nn.Linear( 4096,10)
self.dropout = nn.Dropout(0.25)
self.softmax = nn.LogSoftmax(dim=1)

def forward(self, x):
# add sequence of convolutional and max pooling layers
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = F.relu(self.conv3(x))
x = x.view(x.size(0), 6144 )
nn.Dropout()
x = F.relu(self.fc1(x))
nn.Dropout()
x = F.relu(self.fc2(x))
nn.Dropout()
x = F.relu(self.fc3(x))
x = self.softmax(x)
return x

transform = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,)),
])

model = myModel()

import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

# 打开图片并转成灰度
image = Image.open('9_M.png')
gray=image.convert('L')

# 显示图片
plt.figure("predict")
plt.imshow(gray)
plt.show()

# 转成tensor
tensor = transform(gray)
tensor = tensor.unsqueeze(0)
outputdata = model(inputdata)
ps = torch.exp(outputdata)

top_p, top_class = ps.topk(1, dim=1)

# 输出结果
print(top_p)
print(top_class)


1(0.7305) 3(0.9704) 4(0.8309) 3(0.4391) 8(0.9913)


• 点赞
• 评论 7
• 分享
x

海报分享

扫一扫，分享海报

• 收藏 2
• 手机看

分享到微信朋友圈

x

扫一扫，手机阅读

• 打赏

打赏

yanqs_whu

你的鼓励将是我创作的最大动力

C币 余额
2C币 4C币 6C币 10C币 20C币 50C币
• 一键三连

点赞Mark关注该博主, 随时了解TA的最新博文

05-09 1万+
07-02 361
07-26 1386
04-18 1265
08-12 912
04-19 1362