PyTorch深度学习项目实战：从Notebook到模块化代码的转变

PyTorch深度学习项目实战：从Notebook到模块化代码的转变

pytorch-deep-learning Materials for the Learn PyTorch for Deep Learning: Zero to Mastery course. 项目地址: https://gitcode.com/gh_mirrors/py/pytorch-deep-learning

引言

在深度学习项目开发过程中，我们通常会从Jupyter Notebook开始进行快速原型设计和实验验证。然而，随着项目规模扩大，将代码模块化成为提高可维护性和可重用性的关键步骤。本文将详细介绍如何将PyTorch深度学习项目从Notebook转换为模块化Python脚本。

Notebook与Python脚本的对比

Notebook的优势与局限

Notebook（如Jupyter Notebook或Google Colab）非常适合：

• 快速实验和迭代
• 可视化中间结果
• 交互式调试

但其局限性也很明显：

• 版本控制困难
• 难以重用特定功能
• 代码组织性较差

Python脚本的优势

模块化Python脚本则提供了：

• 更好的代码组织和重用性
• 更完善的版本控制支持
• 更适合生产环境部署
• 更清晰的依赖管理

模块化设计思路

在PyTorch项目中，我们可以将代码按功能划分为多个模块：

1. 数据准备模块 (data_setup.py)
2. 模型构建模块 (model_builder.py)
3. 训练引擎模块 (engine.py)
4. 训练脚本 (train.py)
5. 工具函数模块 (utils.py)

这种模块化设计遵循了单一职责原则，每个模块专注于特定功能。

数据准备模块 (data_setup.py)

数据准备是深度学习项目的第一步，我们将这部分功能封装在data_setup.py中：

"""
功能：创建PyTorch数据加载器用于图像分类任务
"""
import os
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

def create_dataloaders(train_dir, test_dir, transform, batch_size, num_workers=os.cpu_count()):
    """
    创建训练和测试数据加载器
    
    参数:
        train_dir: 训练数据目录路径
        test_dir: 测试数据目录路径  
        transform: 数据预处理变换
        batch_size: 批次大小
        num_workers: 数据加载工作线程数
        
    返回:
        元组(train_dataloader, test_dataloader, class_names)
    """
    # 创建数据集
    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

关键点：

• 使用ImageFolder自动处理图像分类数据集
• 支持多线程数据加载提高效率
• 返回类别名称便于后续使用

模型构建模块 (model_builder.py)

模型定义是深度学习项目的核心，我们将其独立为单独模块：

"""
功能：实现TinyVGG模型架构
"""
import torch
from torch import nn

class TinyVGG(nn.Module):
    """
    TinyVGG模型架构
    
    参数:
        input_shape: 输入通道数
        hidden_units: 隐藏层单元数
        output_shape: 输出类别数
    """
    def __init__(self, input_shape, hidden_units, output_shape):
        super().__init__()
        self.conv_block_1 = nn.Sequential(
            nn.Conv2d(input_shape, hidden_units, 3, 1, 0),
            nn.ReLU(),
            nn.Conv2d(hidden_units, hidden_units, 3, 1, 0),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        
        self.conv_block_2 = nn.Sequential(
            nn.Conv2d(hidden_units, hidden_units, 3, 1, 0),
            nn.ReLU(),
            nn.Conv2d(hidden_units, hidden_units, 3, 1, 0),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(hidden_units*13*13, output_shape)
        )
    
    def forward(self, x):
        x = self.conv_block_1(x)
        x = self.conv_block_2(x)
        x = self.classifier(x)
        return x

模型特点：

• 采用经典的CNN架构设计
• 使用Sequential组织网络层
• 清晰的参数化设计

训练引擎模块 (engine.py)

训练循环是深度学习项目的关键部分：

"""
功能：训练和评估函数
"""
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]:
    """执行单轮训练"""
    model.train()
    train_loss, train_acc = 0, 0
    
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)
        
        # 前向传播
        y_pred = model(X)
        
        # 计算损失
        loss = loss_fn(y_pred, y)
        train_loss += loss.item()
        
        # 优化器清零梯度
        optimizer.zero_grad()
        
        # 反向传播
        loss.backward()
        
        # 参数更新
        optimizer.step()
        
        # 计算准确率
        y_pred_class = torch.argmax(y_pred, dim=1)
        train_acc += (y_pred_class == y).sum().item()/len(y_pred)
    
    train_loss /= len(dataloader)
    train_acc /= len(dataloader)
    
    return train_loss, train_acc

def test_step(model, dataloader, loss_fn, device):
    """执行单轮测试"""
    model.eval()
    test_loss, test_acc = 0, 0
    
    with torch.inference_mode():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            
            # 前向传播
            test_pred = model(X)
            
            # 计算损失和准确率
            loss = loss_fn(test_pred, y)
            test_loss += loss.item()
            
            test_pred_class = torch.argmax(test_pred, dim=1)
            test_acc += (test_pred_class == y).sum().item()/len(test_pred)
    
    test_loss /= len(dataloader)
    test_acc /= len(dataloader)
    
    return test_loss, test_acc

def train(model, train_dataloader, test_dataloader, 
          optimizer, loss_fn, epochs, device):
    """完整训练流程"""
    results = {"train_loss": [], "train_acc": [], 
               "test_loss": [], "test_acc": []}
    
    for epoch in tqdm(range(epochs)):
        train_loss, train_acc = train_step(model, train_dataloader, 
                                         loss_fn, optimizer, device)
        test_loss, test_acc = test_step(model, test_dataloader, 
                                      loss_fn, device)
        
        # 记录结果
        results["train_loss"].append(train_loss)
        results["train_acc"].append(train_acc)
        results["test_loss"].append(test_loss)
        results["test_acc"].append(test_acc)
        
        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}")
    
    return results

关键功能：

• 分离训练和测试步骤
• 支持进度条显示
• 记录并返回训练指标

训练脚本 (train.py)

将各个模块组合成完整的训练流程：

"""
功能：主训练脚本
"""
import os
import torch
import argparse
from torchvision import transforms

# 自定义模块
from going_modular import data_setup, model_builder, engine

# 参数解析
parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str, default="tinyvgg")
parser.add_argument("--batch_size", type=int, default=32)
parser.add_argument("--lr", type=float, default=0.001)
parser.add_argument("--num_epochs", type=int, default=5)
args = parser.parse_args()

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

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor()
])

# 创建数据加载器
train_dir = "data/pizza_steak_sushi/train"
test_dir = "data/pizza_steak_sushi/test"
train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(
    train_dir, test_dir, transform, args.batch_size
)

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

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

# 训练模型
engine.train(
    model=model,
    train_dataloader=train_dataloader,
    test_dataloader=test_dataloader,
    optimizer=optimizer,
    loss_fn=loss_fn,
    epochs=args.num_epochs,
    device=device
)

# 保存模型
model_path = "models/tinyvgg_model.pth"
torch.save(model.state_dict(), model_path)

使用方式：

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

项目目录结构

完整的模块化项目结构如下：

project_root/
├── going_modular/
│   ├── data_setup.py
│   ├── model_builder.py  
│   ├── engine.py
│   ├── train.py
│   └── utils.py
├── data/
│   └── pizza_steak_sushi/
│       ├── train/
│       └── test/
└── models/
    └── tinyvgg_model.pth

最佳实践建议

1. 文档字符串：为每个函数和类添加详细的文档字符串
2. 类型提示：使用Python类型提示提高代码可读性
3. 参数化设计：使函数和类尽可能通用和可配置
4. 错误处理：添加适当的错误处理逻辑
5. 日志记录：使用logging模块记录训练过程

总结

将PyTorch项目从Notebook转换为模块化代码可以带来诸多好处：

• 提高代码可重用性
• 便于团队协作
• 更适合生产环境部署
• 更清晰的代码组织结构

本文介绍的方法可以应用于各种规模的PyTorch深度学习项目，帮助开发者建立更专业、更可维护的代码库。

pytorch-deep-learning Materials for the Learn PyTorch for Deep Learning: Zero to Mastery course. 项目地址: https://gitcode.com/gh_mirrors/py/pytorch-deep-learning