详解用PyTorch搭建一个图像分类神经网络

前言

今天是大年初二，业余时间总结一下在PyTorch框架下，用自组织数据集，自己亲手搭建一个图像分类神经网络的详细步骤.

环境

• 本地编辑器 win 10 + PyCharm
• 本地环境 Anaconda + Python3 + PyTorch
• 训练环境 腾讯云高性能服务器 (斜眼笑~ 勿酸)
• 测试环境 本地

数据集准备

经典的范例是猫狗图像分类，这次不出意外，我们以 Kaggle 上的 Dog vs. Cat 数据集作为本次神经网络训练所需要的数据.

数据集下载地址：https://www.kaggle.com/c/dogs-vs-cats/data

之后放到项目根目录中，命名为train，下分两个文件夹 Cat / Dog.

引入依赖包

from torchvision.datasets import ImageFolder  # 图形数据组织工具
from torchvision import transforms  # 图形变换工具
import torch.nn as nn  # pytorch网络核心依赖
import torch.nn.functional as F # pytorch函数工具箱
import torch.optim as optim  # pytorch优化工具箱
import os  # os文件系统工具包
import torch  # torch核心以依赖

图形变换函数和图形数据加载

# 定义图像变换
tfs = transforms.Compose([
    transforms.Resize((256, 256)), # 规定图形大小
    transforms.ToTensor(),	# 转化为张量数据
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])	# 数据归中
])

# 加载图像数据集，并在加载的时候对图像施加变换
train_data = ImageFolder('./train', tfs)
# print(train_data.class_to_idx)  # {'Cat': 0, 'Dog': 1}
# print(train_data.classes)  # ['Cat', 'Dog']

定义神经网络

class Net(nn.Module):
	# init定义网络中的结构
    def __init__(self):
        super(Net, self).__init__()
        # 3输入，16输出，卷积核(7, 7)，膨胀系数为2
        self.conv1 = nn.Conv2d(3, 16, kernel_size=7, dilation=2)  
        self.conv2 = nn.Conv2d(16, 32, kernel_size=5)
        # dropout
        self.conv2_drop = nn.Dropout2d()
        # 全连接层
        self.fc1 = nn.Linear(55696*2, 1000)
        self.fc2 = nn.Linear(1000, 50)
        self.fc3 = nn.Linear(50, 2)
	
	# forward定义数据在网络中的流向
    def forward(self, x):
    	# 卷积之后做一个最大池化，然后RELU激活
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        # 整形
        x = x.view(-1, 55696*2)
        x = F.relu(self.fc1(x))
        x = F.dropout(x)
        x = F.relu(self.fc2(x))
        x = F.dropout(x)
        x = self.fc3(x)
        return F.log_softmax(x, dim=1)

一个epoch的训练过程

def train_epoch(epoch, model, device, data_loader, optimizer):
	# model进入训练模式
    model.train()
    pid = os.getpid()
    # 加载训练数据，(data, target)二元组表示 “训练数据-训练标签” 
    for batch_idx, (data, target) in enumerate(data_loader):
        # 优化器梯度清 0
        optimizer.zero_grad()
        # 输出特征预测值
        output = model(data.to(device))
        # 计算损失
        loss = F.nll_loss(output, target.to(device))
        # 计算梯度
        loss.backward()
        # 更新梯度
        optimizer.step()
        # 每十个批次打印一下信息
        if batch_idx % 10 == 0:
            print('{}\tTrain Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
            	pid, epoch, batch_idx * len(data), len(data_loader.dataset),
                100. * batch_idx / len(data_loader), loss.item()))

整体训练函数

def train(args, model, device, dataloader_kwargs):
    # 指定随机值，让每次网络的初始值虽然是随机但却是固定的
    torch.manual_seed(args['seed'])
    # 训练数据加载器 (suffle指定随机打乱数据)
    train_loader = torch.utils.data.DataLoader(train_data, batch_size=args['batch_size'], shuffle=True, **dataloader_kwargs)
    # 优化器
    optimizer = optim.SGD(model.parameters(), lr=args['lr'], momentum=args['momentum'])
    # 迭代训练
    for epoch in range(1, args['epochs'] + 1):
        train_epoch(args['epochs'], model, device, train_loader, optimizer)

主函数调用

# 定义参数
# batch_size：每次训练的样本量
# epochs：训练轮次
# lr：学习率
# momentum：动量
# seed：
# log_interval：
args = {
    'batch_size': 32,
    'test_batch_size': 1000,
    'epochs': 10,
    'lr': 0.01,
    'momentum': 0.9,
    'seed': 1,
    'log_interval': 10
}
# 用cuda
use_cuda = True if torch.cuda.is_available() else False

if __name__ == '__main__':
    # 指定训练设备
    device = torch.device("cpu" if use_cuda else "cpu")
    dataloader_kwargs = {'pin_memory': True} if use_cuda else {}
    # 将模型送达设备
    model = Net().to(device)
    # 训练
    train(args, model, device, dataloader_kwargs)

服务器训练结果

已经从服务器下载到本地了

模型测试

先上结果：0.8813643104300544，看来还可以接受.

测试所用的数据集是：https://www.kaggle.com/tongpython/cat-and-dog

def test(model, device):
	# 加载测试数据集
    test_data = ImageFolder(r'./test', tfs)
    test_loader = torch.utils.data.DataLoader(test_data, batch_size=64, shuffle=True)
    lenn = len(test_data)
	# 测试过程不记录权重
    with torch.no_grad():
        acc = 0
        for item in test_loader:
            img, tag = item
            outputs = model(img.to(device))
            # [1] 维度是预测的类别
            predict_y = torch.max(outputs, dim=1)[1]
            # 统计正确分类的个数
            acc += torch.eq(predict_y, tag.to(device)).sum().item()
        acc /= lenn
        print(acc)

主函数改为测试函数

device = torch.device("cpu" if use_cuda else "cpu")
model = Net().to(device)
# 加载模型参数
model.load_state_dict(torch.load('./model.pth', map_location=device), strict=False)
# 测试
test(model, device)

模型预测

def predict(model, img):
    with torch.no_grad():
        out = model(img)
        _, pre = torch.max(out.data, 1)
        return pre.item()

主函数改为：

device = torch.device("cpu" if use_cuda else "cpu")
model = Net().to(device)
model.load_state_dict(torch.load('./model.pth', map_location=device), strict=False)
print(predict(model, tfs(Image.open(r'xxx.jpg')).unsqueeze(0)))

参考和推荐阅读资料

• 深度学习:深入浅出地理解神经网络 https://zhuanlan.zhihu.com/p/61492295
• BP神经网络算法推导及代码实现笔记 https://zhuanlan.zhihu.com/p/38006693
• Pytorch 模型构建、训练、测试及预测 https://blog.csdn.net/andyL_05/article/details/103363603
• 深度学习中Dropout原理解析 https://blog.csdn.net/program_developer/article/details/80737724
• 常用激活函数总结 https://zhuanlan.zhihu.com/p/74393124
• 从 SGD 到 Adam —— 6大常见优化算法总结 https://zhuanlan.zhihu.com/p/64113429