? 卷積神經網路手寫數字影像識別

1、? 匯入相關庫

# pytorch框架
import torch
from torch import nn,optim
from torchvision import datasets,transforms
from torch.utils.data import DataLoader

# 加載圖片
from PIL import Image
from torchvision import models

2、? 匯入手寫資料集

# 從指定路徑加載資料集，不下載，并將其轉化為Tensor格式
train_dataset = datasets.MNIST(root='./',
                train=True,
                transform=transforms.ToTensor(),
                download=False)

test_dataset = datasets.MNIST(root='./',
               train=False,
               transform=transforms.ToTensor(),
               download=False)

3、? 定義資料包裝器

# 每批的資料為64個，用于資料加載，shuffle為打亂資料
batch_size=64
train_loader=DataLoader(dataset=train_dataset,
                       batch_size=batch_size,
                       shuffle=True)
test_loader=DataLoader(dataset=test_dataset,
                      batch_size=batch_size,
                      shuffle=True)

4、? 查看資料維度

# 該維度表示64張圖片，每張的channel為1，寬高為28，28
for i, data in enumerate(train_loader):
    inputs, labels = data
    print(inputs.shape)
    print(labels.shape)
    break

5、? 定義卷積網路層

# nn.Module
class CNN(nn.Module):
	# 初始化相關網路層
    def __init__(self):
    	# 初始化父類
        super(CNN,self).__init__()
        # 定義第一個卷積層
        self.conv1=nn.Sequential(nn.Conv2d(1,64,5,1,2),nn.ReLU(),nn.MaxPool2d(2,2))
        # 定義第二個卷積層
        self.conv2=nn.Sequential(nn.Conv2d(64,128,5,1,2),nn.ReLU(),nn.MaxPool2d(2,2))
        # 定義全連接層
        self.fc1=nn.Sequential(nn.Linear(128*7*7,1000),nn.Dropout(p=0.2),nn.ReLU())
        self.fc2=nn.Sequential(nn.Linear(1000,10),nn.Softmax(dim=1))
    # 定義傳播函式
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.reshape(x.shape[0], -1)
        x = self.fc1(x)
        x = self.fc2(x)
        return x

6、? 定義模型與損失函式、優化器

# 學習率
LR=0.0003
# 定義模型
model=CNN()
# 定義損失函式：交叉熵
loss_fn=nn.CrossEntropyLoss()
# 定義優化器
optimizer=optim.Adam(model.parameters(),LR)

7、? 訓練、測驗函式

def train():
    model.train()
    for i,data in enumerate(train_loader):
        x,y=data
        # 預測結果
        y_pred=model(x)
        # 計算損失
        loss=loss_fn(y_pred,y)
        # 梯度清0
        optimizer.zero_grad()
        # 計算梯度
        loss.backward()
        # 修改權值
        optimizer.step()
def test():
    model.eval()
    correct=0
    for i,data in enumerate(train_loader):
        x,y=data
        y_pred=model(x)
        _,pred=torch.max(y_pred,1)
        correct+=(pred==y).sum()
    print('準確率：',correct.item()/len(train_dataset))
    correct = 0
    for i, data in enumerate(test_loader):
        x,y=data
        y_pred=model(x)
        _,pred=torch.max(y_pred,1)
        correct+=(pred==y).sum()
    print('準確率：',correct.item()/len(test_dataset))

8、? 網路迭代訓練

# 訓練10次模型，不斷修改權重
for epoch in range(1,11):
    print('epoch:',epoch)
    train()
    test()

9、? 保存模型

# 將模型的引數保存到D盤
torch.save(model.state_dict(),r'D:\res.pth')

10、? 加載模型

# 重新定義模型，加載保存的引數
model=CNN()
model.load_state_dict(torch.load(r'D:\res.pth'))
loss_fn=nn.CrossEntropyLoss()
optimizer=optim.Adam(model.parameters(),LR)

11、? 模型測驗

# 將圖片轉化為Tensor
transform = transforms.Compose([
    transforms.ToTensor() 
])
# 打開圖片
img=Image.open(r'C:\Users\HUAWEI\Desktop\7.jpg')
# 改變圖片大小
img=img.resize((28,28))
# 圖片灰度化
img=img.convert('L')
# 將圖片增加一個維度，卷積層需要4維
img=transform(img).unsqueeze(0)
out=model(img)
_,pred=torch.max(out,1)
print('該數字為:',pred.item())