PyTorch 项目模板

对于任何深度学习项目来说，一个简单且设计良好的结构都是至关重要的。经过大量的实践以及参与多个 PyTorch 项目的贡献后，我整理出了一份结合了简洁性、最佳的文件夹结构实践和良好的面向对象设计的 PyTorch 项目模板。

其核心思想在于，每次开始一个新的 PyTorch 项目时，你都会重复许多相同的工作流程。因此，将这些通用的部分封装起来，可以帮助你在每次启动新项目时，只需专注于核心逻辑（如模型架构、训练流程等）即可。

因此，这里提供了一个简单的 PyTorch 模板，旨在帮助你更快地进入核心开发阶段，从而专注于你的核心任务（模型架构、训练流程等）。

为了减少重复性工作，我们建议使用高层次的库。你可以自己编写一个高层次的库，也可以直接使用一些第三方库，例如 ignite、fastai、mmcv 等。这些库能够帮助你用几行代码就实现紧凑但功能齐全的训练循环。在这里，我们以使用 ignite 训练 MNIST 数据集为例进行说明。

需求

• yacs（Yet Another Configuration System）
• PyTorch（一个开源的深度学习平台）
• ignite（用于帮助在 PyTorch 中训练神经网络的高层次库）

目录

• 简要概述
• 详细说明
• 未来工作
• 贡献
• 致谢

简要概述

简而言之，以下是使用此模板的方法。例如，假设你想实现 ResNet-18 来训练 MNIST 数据集，那么你应该按照以下步骤操作：

• 在 modeling 文件夹中创建一个任意命名的 Python 文件，这里我们命名为 example_model.py。在 modeling/__init__.py 文件中，你可以定义一个名为 build_model 的函数来实例化你的模型：

from .example_model import ResNet18

def build_model(cfg):
    model = ResNet18(cfg.MODEL.NUM_CLASSES)
    return model

• 在 engine 文件夹中，创建模型训练函数和推理函数。在训练函数中，你需要编写训练过程的逻辑；可以借助一些第三方库来减少重复代码。

# 训练函数
def do_train(cfg, model, train_loader, val_loader, optimizer, scheduler, loss_fn):
 """
 实现每个 epoch 的逻辑：
 - 根据配置中的迭代次数循环调用训练步骤
 - 使用 summary 添加所需的各类汇总信息
 """
pass

# 推理函数
def inference(cfg, model, val_loader):
"""
实现推理步骤的逻辑：
- 运行 TensorFlow 会话
- 返回需要总结的各项指标
 """
pass

• 在 tools 文件夹中，创建 train.py 文件。在这个文件中，你需要获取“模型”、“数据加载器”、“优化器”以及配置文件的实例：

# 实例化所需模型
model = build_model(cfg)

# 创建数据生成器
train_loader = make_data_loader(cfg, is_train=True)
val_loader = make_data_loader(cfg, is_train=False)

# 创建模型优化器
optimizer = make_optimizer(cfg, model)

• 将上述所有对象传递给 do_train 函数，并开始训练：

# 在这里训练你的模型
do_train(cfg, model，train_loader，val_loader，optimizer，None，F.cross_entropy)

你可以在 model 和 trainer 文件夹中找到模板文件和一个简单的示例，展示如何快速尝试你的第一个模型。

详细说明

├── config
│    └── defaults.py  - 默认配置文件。
│
│
├── configs  
│    └── train_mnist_softmax.yml  - 针对特定模型或数据集的配置文件。
│ 
│
├── data  
│    └── datasets  - 负责所有数据处理的文件夹。
│    └── transforms  - 负责所有数据增强的预处理文件夹。
│    └── build.py  		   - 用于构建数据加载器的文件。
│    └── collate_batch.py   - 该文件负责将多个样本合并成一个小批量。
│
│
├── engine
│   ├── trainer.py     - 包含训练循环的文件。
│   └── inference.py   - 包含推理过程的文件。
│
│
├── layers              - 该项目自定义层的存放目录。
│   └── conv_layer.py
│
│
├── modeling            - 该项目所有模型的存放目录。
│   └── example_model.py
│
│
├── solver             - 该项目优化器的存放目录。
│   └── build.py
│   └── lr_scheduler.py
│   
│ 
├── tools                - 该项目的训练/测试脚本存放处。
│    └── train_net.py  - 一个完整的训练流程示例。
│ 
│ 
└── utils
│    ├── logger.py
│    └── 其他所需工具
│ 
│ 
└── tests					- 该项目的单元测试存放处。
     ├── test_data_sampler.py

未来工作

贡献

欢迎任何形式的改进与贡献。

致谢

Deep-Learning-Project-Template 快速上手指南

本指南旨在帮助开发者快速基于 Deep-Learning-Project-Template 搭建结构清晰、符合最佳实践的 PyTorch 深度学习项目。该模板通过模块化设计，让你只需关注核心模型架构与训练逻辑，无需重复编写基础代码。

环境准备

在开始之前，请确保你的开发环境满足以下要求：

• 操作系统: Linux / macOS / Windows
• Python: 建议 Python 3.6+
• 核心依赖:
  • PyTorch：深度学习框架
  • yacs：配置管理系统
  • ignite：高级训练库（用于简化训练循环）

国内加速建议： 安装 PyTorch 时，推荐使用清华或中科大镜像源以加快下载速度。例如： pip install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

1. 克隆项目仓库

   git clone https://github.com/your-repo/Deep-Learning-Project-Template.git
   cd Deep-Learning-Project-Template
   

2. 安装依赖包 建议使用虚拟环境（如 conda 或 venv），然后安装所需库：

   pip install yacs pytorch ignite -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   (注：若项目根目录包含 requirements.txt，可直接运行 pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple)

基本使用

本模板的核心思想是将通用逻辑封装，你只需修改特定文件夹即可启动新项目。以下以“使用 ResNet-18 训练 MNIST"为例：

1. 定义模型 (modeling)

在 modeling 文件夹下创建模型文件（如 example_model.py），并在 modeling/__init__.py 中注册构建函数：

# modeling/__init__.py
from .example_model import ResNet18

def build_model(cfg):
    model = ResNet18(cfg.MODEL.NUM_CLASSES)
    return model

2. 编写训练与推理逻辑 (engine)

在 engine 文件夹下定义训练流程。利用 ignite 等库可大幅减少样板代码：

# engine/trainer.py
def do_train(cfg, model, train_loader, val_loader, optimizer, scheduler, loss_fn):
    """
    实现训练周期逻辑：
    - 循环迭代次数
    - 执行训练步
    - 添加需要的监控指标
    """
    pass

# engine/inference.py
def inference(cfg, model, val_loader):
    """
    实现推理逻辑：
    - 运行评估
    - 返回需要汇总的指标
    """
    pass

3. 组装启动脚本 (tools)

在 tools 文件夹下的启动脚本（如 train_net.py）中实例化所需对象并启动训练：

# tools/train_net.py

# 1. 构建模型
model = build_model(cfg)

# 2. 构建数据加载器
train_loader = make_data_loader(cfg, is_train=True)
val_loader = make_data_loader(cfg, is_train=False)

# 3. 构建优化器
optimizer = make_optimizer(cfg, model)

# 4. 开始训练
do_train(cfg, model, train_loader, val_loader, optimizer, None, F.cross_entropy)

4. 配置文件管理

• 默认配置: 修改 config/defaults.py 设置全局默认参数。
• 具体任务配置: 在 configs/ 目录下创建具体的 YAML 文件（如 train_mnist_softmax.yml）来覆盖默认设置，指定数据集路径、模型类型等。

运行训练：

python tools/train_net.py --config-file configs/train_mnist_softmax.yml

通过以上步骤，你可以迅速复用模板中的数据处理、日志记录、模型保存等功能，将精力集中在算法创新上。更多详细目录结构说明请参考项目 In Details 章节。