【研究生学习】Pytorch基本知识——图像识别核心模块实战

本博客主要记录一下对于图像识别的实际任务如何使用Pytorch完成一个项目

卷积网络参数定义

首先需要导入必要的模块：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets,transforms
import matplotlib.pyplot as plt
import numpy as np

其次需要定义一些超参数，并且读取数据，采用DataLoader来迭代数据：

# hyper-parameter
input_size = 28
num_classes = 10
num_epochs = 3
batch_size = 64

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

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=True)

接下来就要构建卷积神经网络模块，一般卷积层、relu层和池化层可以写在一起，且卷积最后的结果还是一个特征图，需要把图转换为向量才能做分类或回归任务：

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(     # 输入大小为(1,28,28)
            nn.Conv2d(
                in_channels=1,  # 灰度图
                out_channels=16,    # 要得到多少特征图
                kernel_size=5,  # 卷积核大小
                stride=1,   # 步长
                padding=2,  # 如果希望卷积后大小跟原来一样，需要设置padding=(kernel_size-1)/2 if stride=1
            ),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),    # 池化操作，输出结果为(16,14,14)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 5, 1, 2),     # 输出(32,14,14)
            nn.ReLU(),
            nn.MaxPool2d(2),    # 输出(32,7,7)
        )
        self.out = nn.Linear(32*7*7, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)   # flatten操作，结果为(batch_size, 32*7*7)
        output = self.out(x)
        return output

定义准确率作为评估标准：

def accurancy(predictions, labels):
    pred = torch.max(predictions.data, 1)[1]
    rights = pred.eq(labels.data.view_as(pred)).sum()
    return rights, len(labels)

接下来就可以训练网络：

net = CNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)

for epoch in range(num_epochs):
    train_right = []
    for batch_idx, (data, target) in enumerate(train_loader):
        net.train()
        output = net(data)
        loss = criterion(output, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        right = accurancy(output, target)
        train_right.append(right)

        if batch_idx%100 == 0:
            net.eval()
            val_rights = []
            for (data, target) in test_loader:
                output = net(data)
                right = accurancy(output, target)
                val_rights.append(right)

            # 准确率计算
            train_r = (sum([tup[0] for tup in train_right]), sum([tup[1] for tup in train_right]))
            val_r = (sum([tup[0] for tup in val_rights]), sum([tup[1] for tup in val_rights]))
            print('当前epoch:{} [{}/{} ({:.0f}%)]\t损失:{:.6f}\t训练集准确率:{:.2f}\t测试集正确率:{:.2f}%'.format(
                epoch, batch_idx*batch_size, len(train_loader.dataset),
                100.*batch_idx/len(train_loader),
                loss.data,
                100.*train_r[0].numpy()/train_r[1],
                100.*val_r[0].numpy()/val_r[1]
            ))

部分运行结果如下图所示：