计算机视觉101：5-PyTorch 模块化

计算机视觉101，计算机视觉（Computer Vision, CV）是人工智能的核心领域之一，它赋予机器“看”和“理解”世界的能力，广泛应用于自动驾驶、医疗影像、工业检测、增强现实等前沿场景。然而，从理论到落地，CV工程师不仅需要掌握算法原理，更要具备工程化思维，才能让模型真正服务于现实需求。本课程以​​“学以致用”​​为核心理念，采用​​“理论+代码+部署”​​三位一体的教学方式，带您系统掌握计算机视觉的核心技术栈：

1. 从基础开始​​：理解图像处理（OpenCV）、特征工程（SIFT/SURF）等传统方法，夯实CV基本功
2. 进阶深度学习​​：掌握CNN、Transformer等现代模型（PyTorch），学会数据增强、模型调优技巧
3. 实战部署落地​​：学习模型量化、ONNX转换、边缘计算（TensorRT/RKNN）

让算法在真实场景中高效运行通过​​Python代码示范+工业级案例（如缺陷检测、人脸识别）​​，您将逐步构建从​​算法开发到嵌入式部署​​的完整能力链，最终独立完成一个可落地的CV项目。

无论您是希望转行AI的开发者、在校学生，还是寻求技术突破的工程师，本课程都将助您跨越“纸上谈兵”的瓶颈，成为兼具​​算法能力与工程思维​​的CV实战派。

​​🚀 现在开始，用代码让机器真正“看见”世界！​

什么是模块化？

模块化是指将笔记本代码（来自 Jupyter Notebook 或 Google Colab 笔记本）转换为一系列不同的 Python 脚本，以提供类似的功能。

例如，我们可以将笔记本中的单元格代码转换为以下 Python 文件：

• data_setup.py - 用于准备和下载数据的文件（如果需要的话）。
• engine.py - 包含各种训练函数的文件。
• model_builder.py 或 model.py - 用于创建 PyTorch 模型的文件。
• train.py - 利用其他所有文件并训练目标 PyTorch 模型的文件。
• utils.py - 专门用于实用工具函数的文件。

注意： 上述文件的命名和布局取决于您的使用场景和代码需求。Python 脚本与单个笔记本单元一样通用，这意味着您可以为几乎任何功能创建一个脚本。

示例结构

以下是模块化后的文件结构示例：

project/
│
├── data_setup.py       # 数据准备和下载
├── engine.py           # 训练引擎函数
├── model_builder.py    # 模型构建逻辑
├── train.py            # 主训练入口
└── utils.py            # 工具函数集合

通过这种方式，可以更清晰地组织代码，并提高代码的可维护性和复用性。

为什么需要模块化？

笔记本（Notebooks）非常适合用于迭代探索和快速运行实验。

然而，在更大规模的项目中，你可能会发现 Python 脚本更易于复现且更容易运行。

尽管这是一个有争议的话题，但像 Netflix 这样的公司已经展示了如何在生产环境中使用笔记本作为代码工具。

生产代码 是指运行以向某人或某物提供服务的代码。

例如，如果你有一个在线运行的应用程序，其他人可以访问和使用它，那么运行该应用程序的代码就被认为是 生产代码。

并且像 fast.ai 的 nb-dev（即 notebook development 的缩写）这样的库允许你在 Jupyter Notebook 中编写整个 Python 库（包括文档）。

---

笔记本与 Python 脚本的优缺点

双方都有各自的论据。

以下列表总结了一些主要话题：

优点	缺点
笔记本	易于实验/入门
易于分享（例如：Google Colab 链接）	难以进行版本控制
非常直观	文字和图形可能干扰代码

优点	缺点
Python 脚本	可以将代码打包在一起（避免在不同笔记本之间重复编写相似代码）
可以使用 git 进行版本控制	实验不如笔记本直观（通常需要运行整个脚本而不是单个单元格）
许多开源项目使用脚本
更大的项目可以在云供应商上运行（对笔记本的支持较少）

---

我的工作流程

我通常会在 Jupyter/Google Colab 笔记本中开始机器学习项目，以便快速进行实验和可视化。

然后当我有了一个可行的方案时，我会将最有用的代码片段移动到 Python 脚本中。

编写机器学习代码有许多可能的工作流程。有些人喜欢从脚本开始，另一些人（比如我）则喜欢先从笔记本开始，然后再转向脚本。

---

PyTorch 在实际中的应用

在你的探索过程中，你会看到许多基于 PyTorch 的机器学习项目的代码仓库都提供了如何通过 Python 脚本运行 PyTorch 代码的说明。

例如，你可能会被要求在终端/命令行中运行类似以下的代码来训练模型：

python train.py --model MODEL_NAME --batch_size BATCH_SIZE --lr LEARNING_RATE --num_epochs NUM_EPOCHS

在命令行中运行带有各种超参数设置的 PyTorch train.py 脚本。

在这种情况下，train.py 是目标 Python 脚本，它可能包含用于训练 PyTorch 模型的函数。

而 --model、--batch_size、--lr 和 --num_epochs 被称为参数标志。

你可以将这些参数设置为你喜欢的任何值，只要它们与 train.py 兼容，就会正常工作；如果不兼容，则会报错。

例如，假设我们想用批量大小为 32、学习率为 0.001、训练 10 个 epoch 来训练我们在第 04 笔记本中的 TinyVGG 模型：

python train.py --model tinyvgg --batch_size 32 --lr 0.001 --num_epochs 10

你可以在 train.py 脚本中设置任意数量的参数标志以满足需求。

PyTorch 博客关于训练最先进的计算机视觉模型的文章也采用了这种风格。

我们将要涵盖的内容

本节的主要概念是：将有用的笔记本代码单元转换为可复用的 Python 文件。

这样做可以避免我们重复编写相同的代码。

本节包含两个笔记本：

1. 05. 模块化：第一部分（单元模式） - 该笔记本以传统 Jupyter Notebook/Google Colab 笔记本形式运行，是第 04 笔记本的精简版本。
2. 05. 模块化：第二部分（脚本模式） - 该笔记本与第一个相同，但增加了将每个主要部分转换为 Python 脚本的功能，例如 data_setup.py 和 train.py。

本文档中的文本重点在于代码单元 05. 模块化：第二部分（脚本模式），即顶部带有 %%writefile ... 的代码单元。

为什么分为两部分？

因为有时候学习某件事情的最佳方式是观察它如何与其他事物不同。

如果你同时运行这两个笔记本，你会看到它们之间的差异，而这些差异正是关键的学习点。

并排运行第 05 节的两个笔记本时，你会发现脚本模式笔记本中额外添加了代码单元，用于将单元模式笔记本中的代码转换为 Python 脚本。

我们的目标

在本节结束时，我们希望实现以下两个目标：

1. 能够通过一行命令行代码训练我们在第 04 笔记本中构建的模型（Food Vision Mini）：python train.py。
2. 构建一个可复用的 Python 脚本目录结构，例如：

going_modular/
├── going_modular/
│   ├── data_setup.py
│   ├── engine.py
│   ├── model_builder.py
│   ├── train.py
│   └── utils.py
├── models/
│   ├── 05_going_modular_cell_mode_tinyvgg_model.pth
│   └── 05_going_modular_script_mode_tinyvgg_model.pth
└── data/
    └── pizza_steak_sushi/
        ├── train/
        │   ├── pizza/
        │   │   ├── image01.jpeg
        │   │   └── ...
        │   ├── steak/
        │   └── sushi/
        └── test/
            ├── pizza/
            ├── steak/
            └── sushi/

需要注意的地方

• 文档字符串 (Docstrings) - 编写可重现且易于理解的代码非常重要。基于这一理念，我们将放入脚本中的每个函数或类都遵循 Google 的 Python 文档字符串风格。

• 脚本顶部的导入语句 - 由于我们将创建的所有 Python 脚本都可以视为独立的小程序，因此所有脚本都需要在开头导入所需的模块。例如：

  # 为 train.py 导入所需模块
  
  import os
  import torch
  import data_setup
  import engine
  import model_builder
  import utils
  
  from torchvision import transforms
  

通过这种方式，我们可以确保代码的清晰性和模块化设计。

0. 单元模式与脚本模式

单元模式笔记本

像 05. 模块化第一部分（单元模式） 这样的单元模式笔记本是以常规方式运行的，笔记本中的每个单元格要么是代码，要么是 Markdown。

脚本模式笔记本

像 05. 模块化第二部分（脚本模式） 这样的脚本模式笔记本与单元模式笔记本非常相似，但其中许多代码单元格可能已被转换为 Python 脚本。

注意： 您不需要通过笔记本创建 Python 脚本，也可以直接通过集成开发环境（IDE）如 VS Code 创建它们。将脚本模式笔记本包含在本节中只是为了演示从笔记本过渡到 Python 脚本的一种方法。

1. 获取数据

在每个 05 笔记本中获取数据的过程与笔记本 04 中相同。

通过 Python 的 requests 模块调用 GitHub 下载一个 .zip 文件并解压。

import os
import requests
import zipfile
from pathlib import Path

# 设置数据文件夹路径
data_path = Path("data/")
image_path = data_path / "pizza_steak_sushi"

# 如果图像文件夹不存在，则下载并准备它...
if image_path.is_dir():
    print(f"{image_path} 目录已存在。")
else:
    print(f"未找到 {image_path} 目录，正在创建...")
    image_path.mkdir(parents=True, exist_ok=True)

# 下载披萨、牛排、寿司数据
with open(data_path / "pizza_steak_sushi.zip", "wb") as f:
    request = requests.get("https://github.com/mrdbourke/pytorch-deep-learning/raw/main/data/pizza_steak_sushi.zip")
    print("正在下载披萨、牛排、寿司数据...")
    f.write(request.content)

# 解压披萨、牛排、寿司数据
with zipfile.ZipFile(data_path / "pizza_steak_sushi.zip", "r") as zip_ref:
    print("正在解压披萨、牛排、寿司数据...") 
    zip_ref.extractall(image_path)

# 删除压缩文件
os.remove(data_path / "pizza_steak_sushi.zip")

上述代码执行后，将生成一个名为 data 的文件夹，其中包含另一个名为 pizza_steak_sushi 的目录，该目录按照标准的图像分类格式存储了披萨、牛排和寿司的图像。

data/
└── pizza_steak_sushi/
    ├── train/
    │   ├── pizza/
    │   │   ├── train_image01.jpeg
    │   │   ├── train_image02.jpeg
    │   │   └── ...
    │   ├── steak/
    │   │   └── ...
    │   └── sushi/
    │       └── ...
    └── test/
        ├── pizza/
        │   ├── test_image01.jpeg
        │   └── test_image02.jpeg
        ├── steak/
        └── sushi/

数据结构说明

• train/：训练集，包含三个子目录 pizza、steak 和 sushi，分别存储对应类别的图像。
• test/：测试集，同样包含 pizza、steak 和 sushi 三个子目录，用于模型评估。

此数据结构符合常见的图像分类任务要求，便于后续处理和模型训练。

2. 创建数据集和数据加载器 (data_setup.py)

一旦我们获取了数据，就可以将其转换为 PyTorch 的 Dataset 和 DataLoader（一个用于训练数据，另一个用于测试数据）。

我们将创建有用的 Dataset 和 DataLoader 的代码封装到一个名为 create_dataloaders() 的函数中。

并通过以下代码行将其写入文件：%%writefile going_modular/data_setup.py。

data_setup.py

%%writefile going_modular/data_setup.py

"""
包含用于创建 PyTorch DataLoaders 的功能模块，
适用于图像分类数据。
"""

import os

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

NUM_WORKERS = os.cpu_count()

def create_dataloaders(
    train_dir: str, 
    test_dir: str, 
    transform: transforms.Compose, 
    batch_size: int, 
    num_workers: int = NUM_WORKERS
):
    """
    创建训练和测试的 DataLoaders。

    接收训练目录和测试目录路径，并将它们转换为 PyTorch Datasets，
    然后再转换为 PyTorch DataLoaders。

    参数：
        train_dir (str): 训练目录的路径。
        test_dir (str): 测试目录的路径。
        transform (transforms.Compose): 对训练和测试数据执行的 torchvision 转换。
        batch_size (int): 每个 DataLoader 中每批次的样本数量。
        num_workers (int): 每个 DataLoader 的工作线程数，默认为 CPU 核心数。

    返回：
        一个元组 (train_dataloader, test_dataloader, class_names)。
        其中 class_names 是目标类别的列表。

    示例用法：
        train_dataloader, test_dataloader, class_names = \\
            create_dataloaders(train_dir="path/to/train_dir",
                              test_dir="path/to/test_dir",
                              transform=some_transform,
                              batch_size=32,
                              num_workers=4)
    """

    # 使用 ImageFolder 创建数据集
    train_data = datasets.ImageFolder(train_dir, transform=transform)
    test_data = datasets.ImageFolder(test_dir, transform=transform)

    # 获取类别名称
    class_names = train_data.classes

    # 将图像转换为数据加载器
    train_dataloader = DataLoader(
        train_data,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=True,
    )
    test_dataloader = DataLoader(
        test_data,
        batch_size=batch_size,
        shuffle=False,  # 测试数据无需打乱
        num_workers=num_workers,
        pin_memory=True,
    )

    return train_dataloader, test_dataloader, class_names

如果需要创建 DataLoader，我们现在可以使用 data_setup.py 中的函数，如下所示：

# 导入 data_setup.py
from going_modular import data_setup

# 创建训练/测试数据加载器并获取类别名称列表
train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(...)

注意事项

• ImageFolder：该方法会自动根据文件夹结构推断出类别标签。例如，如果训练目录中有两个子文件夹 cat 和 dog，则类别名称将为 ['cat', 'dog']。
• pin_memory=True：在 GPU 上训练时，启用此选项可以加速数据传输。
• num_workers：设置为 os.cpu_count() 可以充分利用多核 CPU 提高性能。

3. 创建模型 (model_builder.py)

在过去的几个笔记本中（notebook 03 和 notebook 04），我们已经多次构建了 TinyVGG 模型。

因此，将模型放入单独的文件中以便能够重复使用是有意义的。

让我们通过 %%writefile going_modular/model_builder.py 将 TinyVGG() 模型类放入脚本中：

model_builder.py

%%writefile going_modular/model_builder.py
"""
包含用于实例化 TinyVGG 模型的 PyTorch 代码。
"""
import torch
from torch import nn 

class TinyVGG(nn.Module):
    """
    创建 TinyVGG 架构。

    在 PyTorch 中复制 CNN 解释器网站上的 TinyVGG 架构。
    查看原始架构：https://poloclub.github.io/cnn-explainer/

    参数：
        input_shape: 整数，表示输入通道的数量。
        hidden_units: 整数，表示层之间的隐藏单元数量。
        output_shape: 整数，表示输出单元的数量。
    """
    def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:
        super().__init__()
        self.conv_block_1 = nn.Sequential(
            nn.Conv2d(in_channels=input_shape, 
                      out_channels=hidden_units, 
                      kernel_size=3, 
                      stride=1, 
                      padding=0),  
            nn.ReLU(),
            nn.Conv2d(in_channels=hidden_units, 
                      out_channels=hidden_units,
                      kernel_size=3,
                      stride=1,
                      padding=0),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2,
                         stride=2)
        )
        self.conv_block_2 = nn.Sequential(
            nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=0),
            nn.ReLU(),
            nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=0),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            # 这个 in_features 的形状从哪里来？
            # 这是因为网络的每一层都会压缩并改变输入数据的形状。
            nn.Linear(in_features=hidden_units*13*13,
                      out_features=output_shape)
        )

    def forward(self, x: torch.Tensor):
        x = self.conv_block_1(x)
        x = self.conv_block_2(x)
        x = self.classifier(x)
        return x
        # return self.classifier(self.conv_block_2(self.conv_block_1(x))) # <- 利用操作融合的优势

现在，我们无需每次都从头编写 TinyVGG 模型，而是可以通过以下方式导入它：

import torch
# 导入 model_builder.py
from going_modular import model_builder
device = "cuda" if torch.cuda.is_available() else "cpu"

# 从 "model_builder.py" 脚本中实例化模型
torch.manual_seed(42)
model = model_builder.TinyVGG(input_shape=3,
                               hidden_units=10, 
                               output_shape=len(class_names)).to(device)

说明

• 模块化设计：通过将模型定义放入单独的文件中，可以提高代码的可重用性和可维护性。
• 参数解释：
  • input_shape: 输入图像的通道数（例如 RGB 图像为 3）。
  • hidden_units: 卷积层中的隐藏单元数量。
  • output_shape: 输出类别数量。
• 模型结构：
  • 包含两个卷积块 (conv_block_1 和 conv_block_2) 和一个分类器 (classifier)。
  • 使用 nn.Sequential 简化了层的堆叠。
• 设备选择：根据是否支持 CUDA，自动选择运行设备（GPU 或 CPU）。

4. 创建 train_step() 和 test_step() 函数以及 train()

在 notebook 04 中，我们编写了几个训练函数：

1. train_step() - 接收模型、DataLoader、损失函数和优化器，并在 DataLoader 上训练模型。
2. test_step() - 接收模型、DataLoader 和损失函数，并在 DataLoader 上评估模型。
3. train() - 将上述两个函数结合，在给定的轮数（epochs）内执行训练和测试，并返回结果字典。

由于这些函数将是模型训练的核心部分，我们可以将它们放入一个名为 engine.py 的 Python 脚本中，使用以下命令创建文件：%%writefile going_modular/engine.py。

以下是 engine.py 的内容：

## engine.py

```python
%%writefile going_modular/engine.py
"""
包含用于训练和测试 PyTorch 模型的函数。
"""
import torch

from tqdm.auto import tqdm
from typing import Dict, List, Tuple

def train_step(model: torch.nn.Module, 
               dataloader: torch.utils.data.DataLoader, 
               loss_fn: torch.nn.Module, 
               optimizer: torch.optim.Optimizer,
               device: torch.device) -> Tuple[float, float]:
    """训练 PyTorch 模型单个 epoch。

    将目标 PyTorch 模型设置为训练模式，然后运行所有必要的训练步骤（前向传播、损失计算、优化器更新）。

    参数:
        model: 要训练的 PyTorch 模型。
        dataloader: 用于训练模型的 DataLoader 实例。
        loss_fn: 要最小化的 PyTorch 损失函数。
        optimizer: 帮助最小化损失函数的 PyTorch 优化器。
        device: 目标设备，用于计算（例如 "cuda" 或 "cpu"）。

    返回:
        训练损失和训练准确率的元组。形式为 (train_loss, train_accuracy)。例如：

        (0.1112, 0.8743)
    """
    # 将模型设置为训练模式
    model.train()

    # 初始化训练损失和训练准确率
    train_loss, train_acc = 0, 0

    # 遍历数据加载器中的批次
    for batch, (X, y) in enumerate(dataloader):
        # 将数据发送到目标设备
        X, y = X.to(device), y.to(device)

        # 1. 前向传播
        y_pred = model(X)

        # 2. 计算并累积损失
        loss = loss_fn(y_pred, y)
        train_loss += loss.item() 

        # 3. 优化器梯度清零
        optimizer.zero_grad()

        # 4. 反向传播
        loss.backward()

        # 5. 优化器更新
        optimizer.step()

        # 计算并累积准确率指标
        y_pred_class = torch.argmax(torch.softmax(y_pred, dim=1), dim=1)
        train_acc += (y_pred_class == y).sum().item() / len(y_pred)

    # 调整指标以获取每批次的平均损失和准确率
    train_loss = train_loss / len(dataloader)
    train_acc = train_acc / len(dataloader)
    return train_loss, train_acc


def test_step(model: torch.nn.Module, 
              dataloader: torch.utils.data.DataLoader, 
              loss_fn: torch.nn.Module,
              device: torch.device) -> Tuple[float, float]:
    """测试 PyTorch 模型单个 epoch。

    将目标 PyTorch 模型设置为“eval”模式，然后在测试数据集上执行前向传播。

    参数:
        model: 要测试的 PyTorch 模型。
        dataloader: 用于测试模型的 DataLoader 实例。
        loss_fn: 用于计算测试数据损失的 PyTorch 损失函数。
        device: 目标设备，用于计算（例如 "cuda" 或 "cpu"）。

    返回:
        测试损失和测试准确率的元组。形式为 (test_loss, test_accuracy)。例如：

        (0.0223, 0.8985)
    """
    # 将模型设置为评估模式
    model.eval() 

    # 初始化测试损失和测试准确率
    test_loss, test_acc = 0, 0

    # 启用推理上下文管理器
    with torch.inference_mode():
        # 遍历 DataLoader 的批次
        for batch, (X, y) in enumerate(dataloader):
            # 将数据发送到目标设备
            X, y = X.to(device), y.to(device)

            # 1. 前向传播
            test_pred_logits = model(X)

            # 2. 计算并累积损失
            loss = loss_fn(test_pred_logits, y)
            test_loss += loss.item()

            # 计算并累积准确率
            test_pred_labels = test_pred_logits.argmax(dim=1)
            test_acc += ((test_pred_labels == y).sum().item() / len(test_pred_labels))

    # 调整指标以获取每批次的平均损失和准确率
    test_loss = test_loss / len(dataloader)
    test_acc = test_acc / len(dataloader)
    return test_loss, test_acc


def train(model: torch.nn.Module, 
          train_dataloader: torch.utils.data.DataLoader, 
          test_dataloader: torch.utils.data.DataLoader, 
          optimizer: torch.optim.Optimizer,
          loss_fn: torch.nn.Module,
          epochs: int,
          device: torch.device) -> Dict[str, List]:
    """训练和测试 PyTorch 模型。

    在指定的轮数内，通过调用 `train_step()` 和 `test_step()` 函数对目标 PyTorch 模型进行训练和测试，同时在同一轮次循环中完成训练和测试。

    计算、打印并存储评估指标。

    参数:
        model: 要训练和测试的 PyTorch 模型。
        train_dataloader: 用于训练模型的 DataLoader 实例。
        test_dataloader: 用于测试模型的 DataLoader 实例。
        optimizer: 帮助最小化损失函数的 PyTorch 优化器。
        loss_fn: 用于计算两个数据集损失的 PyTorch 损失函数。
        epochs: 表示训练多少轮的整数。
        device: 目标设备，用于计算（例如 "cuda" 或 "cpu"）。

    返回:
        包含训练和测试损失以及训练和测试准确率指标的字典。每个指标的值存储在一个列表中，对应每个 epoch。
        形式为：{train_loss: [...],
                 train_acc: [...],
                 test_loss: [...],
                 test_acc: [...]}

        例如，如果训练 epochs=2：
        {train_loss: [2.0616, 1.0537],
         train_acc: [0.3945, 0.3945],
         test_loss: [1.2641, 1.5706],
         test_acc: [0.3400, 0.2973]}
    """
    # 创建空的结果字典
    results = {"train_loss": [],
               "train_acc": [],
               "test_loss": [],
               "test_acc": []
              }

    # 在多个 epoch 中循环执行训练和测试步骤
    for epoch in tqdm(range(epochs)):
        train_loss, train_acc = train_step(model=model,
                                          dataloader=train_dataloader,
                                          loss_fn=loss_fn,
                                          optimizer=optimizer,
                                          device=device)
        test_loss, test_acc = test_step(model=model,
                                        dataloader=test_dataloader,
                                        loss_fn=loss_fn,
                                        device=device)

        # 打印当前进度
        print(
            f"Epoch: {epoch+1} | "
            f"train_loss: {train_loss:.4f} | "
            f"train_acc: {train_acc:.4f} | "
            f"test_loss: {test_loss:.4f} | "
            f"test_acc: {test_acc:.4f}"
        )

        # 更新结果字典
        results["train_loss"].append(train_loss)
        results["train_acc"].append(train_acc)
        results["test_loss"].append(test_loss)
        results["test_acc"].append(test_acc)

    # 返回填充后的结果字典
    return results

现在我们有了 engine.py 脚本，可以通过以下方式导入其中的函数：

# 导入 engine.py
from going_modular import engine

# 调用 engine.py 中的 train() 函数
engine.train(...)

5. 创建用于保存模型的函数 (utils.py)

在训练模型的过程中，或者训练完成后，你可能希望保存模型。

由于我们在之前的笔记本中已经多次编写了保存模型的代码，因此将其转换为一个函数并保存到文件中是合理的。

通常，我们会将辅助函数存储在一个名为 utils.py 的文件中（代表 utilities，即工具）。

让我们通过以下命令将 save_model() 函数保存到名为 utils.py 的文件中：

utils.py

    %%writefile going_modular/utils.py
    """
    包含各种用于 PyTorch 模型训练和保存的实用函数。
    """
    import torch
    from pathlib import Path

    def save_model(model: torch.nn.Module,
                   target_dir: str,
                   model_name: str):
        """将 PyTorch 模型保存到目标目录。

        参数:
            model: 要保存的目标 PyTorch 模型。
            target_dir: 保存模型的目标目录。
            model_name: 保存模型的文件名。应包含 ".pth" 或 ".pt" 作为文件扩展名。

        示例用法:
            save_model(model=model_0,
                       target_dir="models",
                       model_name="05_going_modular_tingvgg_model.pth")
        """
        # 创建目标目录
        target_dir_path = Path(target_dir)
        target_dir_path.mkdir(parents=True, exist_ok=True)

        # 创建模型保存路径
        assert model_name.endswith(".pth") or model_name.endswith(".pt"), "model_name 应以 '.pt' 或 '.pth' 结尾"
        model_save_path = target_dir_path / model_name

        # 保存模型的状态字典
        print(f"[INFO] 正在保存模型到: {model_save_path}")
        torch.save(obj=model.state_dict(), f=model_save_path)

现在，如果我们想使用 save_model() 函数，无需重新编写它，只需导入并调用即可：

# 导入 utils.py
from going_modular import utils

# 将模型保存到文件
utils.save_model(
    model=..., 
    target_dir=...,
    model_name=...
)

格式优化后的 Markdown 内容

## 5. 创建用于保存模型的函数 (`utils.py`)

在训练模型的过程中，或者训练完成后，你可能希望保存模型。

由于我们在之前的笔记本中已经多次编写了保存模型的代码，因此将其转换为一个函数并保存到文件中是合理的。

通常，我们会将辅助函数存储在一个名为 **`utils.py`** 的文件中（代表 utilities，即工具）。

让我们通过以下命令将 `save_model()` 函数保存到名为 `utils.py` 的文件中：

### `utils.py` 文件内容

```python
%%writefile going_modular/utils.py
"""
包含各种用于 PyTorch 模型训练和保存的实用函数。
"""
import torch
from pathlib import Path

def save_model(model: torch.nn.Module, target_dir: str, model_name: str):
    """将 PyTorch 模型保存到目标目录。

    参数:
        model: 要保存的目标 PyTorch 模型。
        target_dir: 保存模型的目标目录。
        model_name: 保存模型的文件名。应包含 `.pth` 或 `.pt` 作为文件扩展名。

    示例用法:
        save_model(model=model_0,
                   target_dir="models",
                   model_name="05_going_modular_tingvgg_model.pth")
    """
    # 创建目标目录
    target_dir_path = Path(target_dir)
    target_dir_path.mkdir(parents=True, exist_ok=True)

    # 创建模型保存路径
    assert model_name.endswith(".pth") or model_name.endswith(".pt"), "model_name 应以 '.pt' 或 '.pth' 结尾"
    model_save_path = target_dir_path / model_name

    # 保存模型的状态字典
    print(f"[INFO] 正在保存模型到: {model_save_path}")
    torch.save(obj=model.state_dict(), f=model_save_path)

使用方法

如果需要使用 save_model() 函数，可以通过以下方式导入并调用：

# 导入 utils.py
from going_modular import utils

# 将模型保存到文件
utils.save_model(
    model=..., 
    target_dir=...,
    model_name=...
)

这样可以避免重复编写代码，并提高代码的可维护性。

6. 训练、评估并保存模型 (train.py)

正如之前讨论的，你经常会遇到将所有功能整合到一个 train.py 文件中的 PyTorch 仓库。

这个文件本质上是在说“使用可用的数据训练模型”。

在我们的 train.py 文件中，我们将把其他 Python 脚本的功能结合起来，并用它来训练一个模型。

这样我们就可以通过命令行的一行代码来训练一个 PyTorch 模型：

python train.py

为了创建 train.py，我们将按照以下步骤进行：

1. 导入各种依赖项，主要包括 torch、os、torchvision.transforms 以及来自 going_modular 目录的所有脚本：data_setup、engine、model_builder 和 utils。
2. 注意：由于 train.py 将位于 going_modular 目录内，因此我们可以直接通过 import ... 导入其他模块，而无需使用 from going_modular import ...。
3. 设置各种超参数，例如批量大小（batch size）、训练轮数（epochs）、学习率（learning rate）和隐藏单元数（hidden units）。这些超参数未来可以通过 Python 的 argparse 模块设置。
4. 设置训练和测试数据目录。
5. 编写与设备无关的代码（device-agnostic code）。
6. 创建必要的数据转换（transforms）。
7. 使用 data_setup.py 创建数据加载器（DataLoaders）。
8. 使用 model_builder.py 创建模型。
9. 设置损失函数和优化器。
10. 使用 engine.py 训练模型。
11. 使用 utils.py 保存模型。

我们可以通过在 notebook 单元格中运行以下命令来创建该文件：

%%writefile going_modular/train.py

以下是完整的 train.py 文件内容：

%%writefile going_modular/train.py
"""
使用与设备无关的代码训练 PyTorch 图像分类模型。
"""

import os
import torch
import data_setup
import engine
import model_builder
import utils

from torchvision import transforms

# 设置超参数
NUM_EPOCHS = 5
BATCH_SIZE = 32
HIDDEN_UNITS = 10
LEARNING_RATE = 0.001

# 设置目录
train_dir = "data/pizza_steak_sushi/train"
test_dir = "data/pizza_steak_sushi/test"

# 设置目标设备
device = "cuda" if torch.cuda.is_available() else "cpu"

# 创建数据转换
data_transform = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor()
])

# 使用 data_setup.py 创建 DataLoaders
train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(
    train_dir=train_dir,
    test_dir=test_dir,
    transform=data_transform,
    batch_size=BATCH_SIZE
)

# 使用 model_builder.py 创建模型
model = model_builder.TinyVGG(
    input_shape=3,
    hidden_units=HIDDEN_UNITS,
    output_shape=len(class_names)
).to(device)

# 设置损失函数和优化器
loss_fn = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)

# 使用 engine.py 开始训练
engine.train(
    model=model,
    train_dataloader=train_dataloader,
    test_dataloader=test_dataloader,
    loss_fn=loss_fn,
    optimizer=optimizer,
    epochs=NUM_EPOCHS,
    device=device
)

# 使用 utils.py 保存模型
utils.save_model(
    model=model,
    target_dir="models",
    model_name="05_going_modular_script_mode_tinyvgg_model.pth"
)

太棒了！

现在我们可以通过在命令行中运行以下命令来训练一个 PyTorch 模型：

python train.py

这样做将利用我们创建的所有其他代码脚本。

如果需要，我们还可以调整 train.py 文件以使用 Python 的 argparse 模块接受命令行参数输入。这将允许我们提供不同的超参数设置，例如：

python train.py --model MODEL_NAME --batch_size BATCH_SIZE --lr LEARNING_RATE --num_epochs NUM_EPOCHS

这样可以更灵活地控制训练过程。

课外学习

• 若要了解有关构建 Python 项目的更多信息，请查阅 Real Python 关于 Python 应用程序布局 的指南。
• 若要获取关于 PyTorch 代码样式的灵感，请参考 Igor Susmelj 编写的 PyTorch 样式指南（本章中的许多样式规则基于此指南以及各种类似的 PyTorch 仓库）。
• 若要查看由 PyTorch 团队编写的示例 train.py 脚本以及其他用于训练最先进的图像分类模型的 PyTorch 脚本，请访问他们在 GitHub 上的 classification 仓库。

---