Pytorch Tutorial

全栈O-Jay

已于 2024-06-03 14:53:39 修改

阅读量506

点赞数 5

分类专栏：人工智能 Python 文章标签： pytorch 人工智能 python 深度学习机器学习

于 2024-03-12 16:18:14 首次发布

本文链接：https://blog.csdn.net/weixin_45116099/article/details/136651651

版权

人工智能同时被 2 个专栏收录

57 篇文章 31 订阅

订阅专栏

Python

35 篇文章 7 订阅

订阅专栏

本文详细介绍了如何使用PyTorch库训练VGG-11神经网络，包括数据预处理、模型构建、优化器设置、学习率调度以及在CIFAR-10数据集上的训练和评估过程。还展示了模型的可视化和性能评估。

摘要由CSDN通过智能技术生成

本教程将详细展示如何使用PyTorch训练神经网络，并给出完整代码和关键注释（ipython），建议使用CoLab的GPU来编译代码。

环境初始化

!pip install torchprofile 1>/dev/null 
#torchprofile用于分析PyTorch模型，帮助理解模型的计算复杂度和参数数量等；1>/dev/null表示不显示安装过程中的任何输出信息，为了简洁。

#一堆库
import random
from collections import OrderedDict, defaultdict

import numpy as np
import torch
from matplotlib import pyplot as plt
from torch import nn
from torch.optim import *
from torch.optim.lr_scheduler import *
from torch.utils.data import DataLoader
from torchprofile import profile_macs
from torchvision.datasets import *
from torchvision.transforms import *
from tqdm.auto import tqdm

#定义随机种子，为了能复现
random.seed(0)
np.random.seed(0)
torch.manual_seed(0)

数据准备

本教程用 CIFAR-10 作为训练数据集。该数据集总共有60,000张图片。这些图片被分为50,000张训练图片和10,000张测试图片，包含来自10个类的图像，其中每个图像的大小为3x32x32，即大小为32x32像素的3通道彩色图像。

transforms = {
  "train": Compose([
    RandomCrop(32, padding=4),
    RandomHorizontalFlip(),
    ToTensor(),
  ]),
  "test": ToTensor(),
}

dataset = {}
for split in ["train", "test"]:
  dataset[split] = CIFAR10(
    root="data/cifar10",
    train=(split == "train"),
    download=True,
    transform=transforms[split],
  )

#可视化图像：从测试数据集中为每个类别抽取四个样本，并将这些样本的图像和标签可视化。
samples = [[] for _ in range(10)]  #创建一个列表samples，其中包含10个子列表，对应于数据集中的10个类别。每个子列表用于存储该类别的样本图像
for image, label in dataset["test"]:  #遍历测试数据集dataset["test"]的每一个图像和标签对。对于每个样本，检查其标签对应的子列表中的样本数量。如果某个类别的样本数少于4个，就将当前图像添加到该类别对应的子列表中
  if len(samples[label]) < 4:
    samples[label].append(image)

plt.figure(figsize=(20, 9)) #设置一个适合显示这些图像的图形大小。
for index in range(40):
  label = index % 10 #01234567890123...每行依次显示每个label的图像
  image = samples[label][index // 10] #0000000000111...每列依次显示单个label的图像

  # 图片格式由 CHW 转换到 HWC，为了可视化
  image = image.permute(1, 2, 0)

  # 将类索引转换为类名
  label = dataset["test"].classes[label]

  # 画图 4 * 10
  plt.subplot(4, 10, index + 1)
  plt.imshow(image)
  plt.title(label)
  plt.axis("off")
plt.show()

在这里插入图片描述

为了训练神经网络，我们需要批量输入数据。我们创建批处理大小为512的数据加载器（data loaders）:

dataflow = {}
for split in ['train', 'test']:
  dataflow[split] = DataLoader(
    dataset[split],
    batch_size=512,
    shuffle=(split == 'train'),
    num_workers=0,
    pin_memory=True,
  )

for inputs, targets in dataflow["train"]:
  print("[inputs] dtype: {}, shape: {}".format(inputs.dtype, inputs.shape))
  print("[targets] dtype: {}, shape: {}".format(targets.dtype, targets.shape))
  break

在这里插入图片描述

模型搭建

我们将使用VGG-11的一个变体(具有更少的下样本和更小的分类器)作为我们的模型。

class VGG(nn.Module):
  ARCH = [64, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'] #网络架构：8个卷积层 + 4个池化层

  def __init__(self) -> None:
    super().__init__()

    layers = []
    counts = defaultdict(int)

    def add(name: str, layer: nn.Module) -> None: #用于添加层到layers列表
      layers.append((f"{name}{counts[name]}", layer))
      counts[name] += 1

    in_channels = 3
    for x in self.ARCH:
      if x != 'M':
        # conv-bn-relu
        add("conv", nn.Conv2d(in_channels, x, 3, padding=1, bias=False))
        add("bn", nn.BatchNorm2d(x))
        add("relu", nn.ReLU(True))
        in_channels = x
      else:
        # maxpool
        add("pool", nn.MaxPool2d(2))

    self.backbone = nn.Sequential(OrderedDict(layers)) #将层列表转换为有序字典，并通过nn.Sequential创建一个顺序容器self.backbone，这使得输入数据可以顺序通过定义的所有层。
    self.classifier = nn.Linear(512, 10) #线性层（nn.Linear），用于将卷积网络提取的特征映射到类别标签上。

  def forward(self, x: torch.Tensor) -> torch.Tensor:
    # backbone: [N, 3, 32, 32] => [N, 512, 2, 2]
    x = self.backbone(x)  #特征提取

    # avgpool: [N, 512, 2, 2] => [N, 512]
    x = x.mean([2, 3])  #对特征图进行全局平均池化

    # classifier: [N, 512] => [N, 10]
    x = self.classifier(x) #通过分类器得到最终的类别预测
    return x 

model = VGG().cuda() #创建了VGG类的一个实例，并通过.cuda()方法将模型的所有参数和缓冲区移动到GPU上，以利用GPU加速计算。（假设你的环境支持CUDA）

#详细看一下模型结构
print(model.backbone)

在这里插入图片描述

#详细分析一下模型的参数

#计算模型大小：
num_params = 0
for param in model.parameters():
  if param.requires_grad:
    num_params += param.numel()
print("#Params:", num_params)
#Params: 9228362

#计算模型的计算开销，由multiply–accumulate operations (MACs，乘法累加操作)来衡量
num_macs = profile_macs(model, torch.zeros(1, 3, 32, 32).cuda())
print("#MACs:", num_macs)
#MACs: 606164480

#该模型有9.2M个参数，需要606M次乘法和累加操作进行一次推理。

模型训练

优化器

optimizer = SGD(
  model.parameters(),
  lr=0.4,
  momentum=0.9,
  weight_decay=5e-4, #权重衰减（L2正则化），有助于防止模型过拟合
)

num_epochs = 20
steps_per_epoch = len(dataflow["train"])

# 分段线性学习率调度。学习率随着训练步数的增加先线性增大，达到一定值后再线性减小。
lr_lambda = lambda step: np.interp(
  [step / steps_per_epoch],
  [0, num_epochs * 0.3, num_epochs],
  [0, 1, 0]
)[0]

# Visualize the learning rate schedule
steps = np.arange(steps_per_epoch * num_epochs)
plt.plot(steps, [lr_lambda(step) * 0.4 for step in steps])
plt.xlabel("Number of Steps")
plt.ylabel("Learning Rate")
plt.grid("on")
plt.show()

scheduler = LambdaLR(optimizer, lr_lambda) #应用学习率调度器

在这里插入图片描述

训练和评估函数

def train(
  model: nn.Module,
  dataflow: DataLoader,
  criterion: nn.Module,
  optimizer: Optimizer,
  scheduler: LambdaLR,
) -> None:
  model.train() #告诉PyTorch模型现在处于训练模式，这对于某些特定层如Dropout和BatchNorm是必要的，因为它们在训练和评估时的行为不同

  for inputs, targets in tqdm(dataflow, desc='train', leave=False):
    # Move the data from CPU to GPU
    inputs = inputs.cuda()
    targets = targets.cuda()

    # 在每次的参数更新前，需要将梯度归零，防止梯度在反向传播时累积
    optimizer.zero_grad()

    # Forward inference
    outputs = model(inputs)
    loss = criterion(outputs, targets)

    # Backward propagation
    loss.backward()

    # 根据计算出的梯度更新模型参数
    optimizer.step()
    # 根据学习率调度器更新学习率
    scheduler.step()


@torch.inference_mode() # 禁用梯度计算
def evaluate(
  model: nn.Module,
  dataflow: DataLoader
) -> float:
  model.eval() ##告诉PyTorch模型现在处于评估模式，关闭Dropout和BatchNorm的特定训练行为。

  num_samples = 0
  num_correct = 0

  for inputs, targets in tqdm(dataflow, desc="eval", leave=False):
    # Move the data from CPU to GPU
    inputs = inputs.cuda()
    targets = targets.cuda()

    # Inference
    outputs = model(inputs)

    # 将模型输出（通常是逻辑值或概率）转换为类别索引
    outputs = outputs.argmax(dim=1)

    # Update metrics
    num_samples += targets.size(0)
    num_correct += (outputs == targets).sum()

  return (num_correct / num_samples * 100).item()

开始训练

for epoch_num in tqdm(range(1, num_epochs + 1)):
  train(model, dataflow["train"], criterion, optimizer, scheduler)
  metric = evaluate(model, dataflow["test"])
  print(f"epoch {epoch_num}:", metric)

在这里插入图片描述

可视化

可视化模型的预测，看看模型的真实表现。

plt.figure(figsize=(20, 10))
for index in range(40):
  image, label = dataset["test"][index+66]

  # Model inference
  model.eval()
  with torch.inference_mode():
    pred = model(image.unsqueeze(dim=0).cuda())
    pred = pred.argmax(dim=1)

  # Convert from CHW to HWC for visualization
  image = image.permute(1, 2, 0)

  # Convert from class indices to class names
  pred = dataset["test"].classes[pred]
  label = dataset["test"].classes[label]

  # Visualize the image
  plt.subplot(4, 10, index + 1)
  plt.imshow(image)
  plt.title(f"pred: {pred}" + "\n" + f"label: {label}")
  plt.axis("off")
plt.show()