【小白指南】pytorch demo 时装分类

FashionMNIST数据集 分类任务练手

# 导入包
import os
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

数据获取

上来先加载个数据

加载数据前设置数据的规范化格式

# 首先设置数据变换
from torchvision import transforms

image_size = 28
data_transform = transforms.Compose([
    transforms.ToPILImage(),  
     # 这一步取决于后续的数据读取方式，如果使用内置数据集读取方式则不需要
    transforms.Resize(image_size),
    transforms.ToTensor()
])

## 读取数据集，小白可以先用这种，但这种读取方法直接就能调用结构化的数据，但实际应用中自己的数据格式是多种多样的，需要自定义格式，后面咱们会讲如何自定义dataset

from torchvision import datasets

train_data = datasets.FashionMNIST(root='./', train=True, download=True, transform=data_transform)
test_data = datasets.FashionMNIST(root='./', train=False, download=True, transform=data_transform)

把数据装入dataloader，dataloader就是个按批次自动吐数据的工具，配置好dataloader，模型就能自助取数据了

# 参数配置
batch_size = 256 # 批次大小，一次同时推理多少张图片
num_workers = 4 # 多线程

# 训练集dataloader
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=num_workers, drop_last=True)

# 测试集dataloader
test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False, num_workers=num_workers)

模型构建

模型训练三大件： 模型结构、损失函数、优化器，缺一不可

1、模型

模型结构选用最简单的CNN，模型类的定义如下：

# 定义模型结构
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 32, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, stride=2),
            nn.Dropout(0.3),
            nn.Conv2d(32, 64, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, stride=2),
            nn.Dropout(0.3)
        )
        self.fc = nn.Sequential(
            nn.Linear(64*4*4, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )
        
    def forward(self, x):
        x = self.conv(x)
        x = x.view(-1, 64*4*4)
        x = self.fc(x)
        # x = nn.functional.normalize(x)
        return x

初始化模型

# 指定设备，是cpu还是GPU
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")


# 模型类操作
model = Net() # 模型类初始化
model = model.to(device) # 将模型存到指定设备上，pytorch不手动存就会在CPU上呆着

2、损失函数 和 优化器

torch.nn自带很多损失函数，基本够用，分类任务一般选择交叉熵损失。

ps: 损失函数是用来计算推理值和理想值之间距离的，pytorch会存下损失函数的运算过程，运用链式法则对参数求偏导，偏导值进行优化器计算出参数在此时的梯度，参数值减去梯度就完成了一次优化，求偏导过程从结果向前求，所以是反向传播，完成反向传播和梯度更新，模型的一次训练就完成了，我们管这一整个过程叫一个iter, 将训练集的数据推理一遍叫一个epoch。

ps2: 优化器是在计算梯度时起作用。因为直接用所求偏导的值作为该步的梯度会导致模型很容易错过最优解，梯度下降太小会导致模型陷入局部最优解，梯度太大会越过最优解，或梯度无法下降，损失值出现打滑现象。

# 配置损失函数
criterion = nn.CrossEntropyLoss()

# 设定优化器
lr = 1e-4

# adam的学习率会自动调整，只要给个初始值就可以，但初始一般很小，过大会导致训练效果变差
optimizer = optim.Adam(model.parameters(), lr=0.001)

训练逻辑

定义训练过程函数

# 训练过程逻辑函数
def train(epoch):
    model.train() # 模型进入训练模式
    train_loss = 0 # 损失值初始化

    # 加载数据
    for data, label in train_loader:
        data, label = data.to(device), label.to(device) # 把数据送到设备中
        optimizer.zero_grad() # 梯度清零，不清零会累积上一步的梯度
        output = model(data) # 模型推理，即前向传播
        loss = criterion(output, label) # 计算损失值
        loss.backward() # 反向传播
        optimizer.step() # 参数更新
        train_loss += loss.item()*data.size(0) # 累计损失
    train_loss = train_loss/len(train_loader.dataset) # iter平均损失
    print('Epoch: {} \tTraining Loss: {:.6f}'.format(epoch, train_loss)) # 打印训练结果

定义验证过程函数

# 验证推理函数，验证与训练的区别是不更新梯度，即为只有前向推理，没有反向传播
def val(epoch):       
    model.eval() # 模型进入推理模式
    val_loss = 0
    gt_labels = []
    pred_labels = []
    with torch.no_grad():
        for data, label in test_loader:
            data, label = data.to(device), label.to(device)
            output = model(data)
            
            # 计算出output对应label
            preds = torch.argmax(output, 1)
            gt_labels.append(label.cpu().data.numpy())
            pred_labels.append(preds.cpu().data.numpy())

            # 只计算损失，不更新梯度
            loss = criterion(output, label)
            val_loss += loss.item()*data.size(0)
    val_loss = val_loss/len(test_loader.dataset)
    gt_labels, pred_labels = np.concatenate(gt_labels), np.concatenate(pred_labels)
    acc = np.sum(gt_labels==pred_labels)/len(pred_labels)
    print('Epoch: {} \tValidation Loss: {:.6f}, Accuracy: {:6f}'.format(epoch, val_loss, acc))

开练

epochs = 20 # 训练轮次

# 开始训练
for epoch in range(1, epochs+1):
    train(epoch)
    val(epoch)

结果保存，如果结果满意，记得保存下来

save_path = "./FahionModel.pkl"
torch.save(model, save_path)

致谢：感谢datawhale组织 ZhikangNiu 提供的详细讲义

讲义地址：点击此处