PyTorch 模板项目

让 PyTorch 深度学习项目变得简单。

• PyTorch 模板项目
  • 要求
  • 特性
  • 文件夹结构
  • 使用方法
    • 配置文件格式
    • 使用配置文件
    • 从检查点恢复
  • 使用多 GPU
  • 自定义
    • 自定义 CLI 选项
    • 数据加载器
    • 训练器
    • 模型
    • 损失函数
    • 指标
    • 额外的日志记录
    • 验证数据
    • 检查点
  • TensorBoard 可视化
  • 贡献
  • 待办事项
  • 许可证
  • 致谢

要求

• Python >= 3.5（推荐 3.6）
• PyTorch >= 0.4（推荐 1.2）
• tqdm（test.py 可选）
• tensorboard >= 1.14（参见 TensorBoard 可视化）

特性

• 清晰的文件夹结构，适用于多种深度学习项目。
• 支持 .json 配置文件，方便参数调优。
• 可自定义的命令行选项，进一步简化参数调整。
• 检查点保存与恢复功能。
• 抽象基类加速开发：
  • BaseTrainer 处理检查点的保存与恢复、训练过程日志记录等。
  • BaseDataLoader 负责批次生成、数据打乱及验证集划分。
  • BaseModel 提供基础的模型摘要信息。

文件夹结构

pytorch-template/
│
├── train.py - 启动训练的主要脚本
├── test.py - 评估训练好的模型
│
├── config.json - 存储训练配置
├── parse_config.py - 处理配置文件和命令行选项的类
│
├── new_project.py - 使用模板文件初始化新项目
│
├── base/ - 抽象基类
│   ├── base_data_loader.py
│   ├── base_model.py
│   └── base_trainer.py
│
├── data_loader/ - 所有与数据加载相关的内容放在此处
│   └── data_loaders.py
│
├── data/ - 默认存储输入数据的目录
│
├── model/ - 模型、损失函数和指标
│   ├── model.py
│   ├── metric.py
│   └── loss.py
│
├── saved/
│   ├── models/ - 训练好的模型保存于此
│   └── log/ - TensorBoard 和日志输出的默认目录
│
├── trainer/ - 训练器
│   └── trainer.py
│
├── logger/ - 用于 TensorBoard 可视化和日志记录的模块
│   ├── visualization.py
│   ├── logger.py
│   └── logger_config.json
│  
└── utils/ - 小工具函数
    ├── util.py
    └── ...

使用方法

此仓库中的代码是一个 MNIST 数据集的模板示例。尝试运行 python train.py -c config.json 来执行代码。

配置文件格式

配置文件采用 .json 格式：

{
  "name": "Mnist_LeNet",        // 训练会话名称
  "n_gpu": 1,                   // 用于训练的 GPU 数量。
  
  "arch": {
    "type": "MnistModel",       // 要训练的模型架构名称
    "args": {

    }                
  },
  "data_loader": {
    "type": "MnistDataLoader",         // 选择数据加载器
    "args":{
      "data_dir": "data/",             // 数据集路径
      "batch_size": 64,                // 批次大小
      "shuffle": true,                 // 在分割前打乱训练数据
      "validation_split": 0.1          // 验证集大小。可以是小数比例或样本数量
      "num_workers": 2,                // 用于数据加载的 CPU 进程数
    }
  },
  "optimizer": {
    "type": "Adam",
    "args":{
      "lr": 0.001,                     // 学习率
      "weight_decay": 0,               // （可选）权重衰减
      "amsgrad": true
    }
  },
  "loss": "nll_loss",                  // 损失函数
  "metrics": [
    "accuracy", "top_k_acc"            // 评估指标列表
  ],                         
  "lr_scheduler": {
    "type": "StepLR",                  // 学习率调度器
    "args":{
      "step_size": 50,          
      "gamma": 0.1
    }
  },
  "trainer": {
    "epochs": 100,                     // 训练轮数
    "save_dir": "saved/",              // 检查点保存在 save_dir/models/name 目录下
    "save_freq": 1,                    // 每 save_freq 轮保存一次检查点
    "verbosity": 2,                    // 0: 静默，1: 每轮打印，2: 全部详细信息
  
    "monitor": "min val_loss"          // 模型性能监控模式及指标。设置为 'off' 可禁用。
    "early_stop": 10	                 // 等待提前停止的轮数。设置为 0 可禁用。
  
    "tensorboard": true,               // 启用 TensorBoard 可视化
  }
}

如有需要，可添加其他配置。

使用配置文件

修改 .json 配置文件中的配置项后，运行：

python train.py --config config.json

从检查点恢复

您可以从之前保存的检查点继续训练，方法如下：

python train.py --resume path/to/checkpoint

使用多 GPU

通过将配置文件中的 n_gpu 参数设置为更大的数值，即可启用多 GPU 训练。如果配置的 GPU 数量少于可用数量，则默认使用前 n 个设备。您还可以通过 CUDA 环境变量指定可用 GPU 的索引。

python train.py --device 2,3 -c config.json

这等同于：

CUDA_VISIBLE_DEVICES=2,3 python train.py -c config.py

自定义

项目初始化

使用 new_project.py 脚本，可以创建包含模板文件的新项目目录。例如，运行 python new_project.py ../NewProject 将创建一个名为 NewProject 的新项目文件夹。该脚本会过滤掉不必要的文件，如缓存、Git 文件或 README 文件。

自定义 CLI 选项

修改配置文件中的值是一种干净、安全且简便的超参数调优方式。然而，当某些值需要频繁或快速更改时，使用命令行选项可能更为合适。

此模板默认使用存储在 JSON 文件中的配置，但通过按如下方式注册自定义选项，你可以利用 CLI 标志来修改其中的部分设置。

# 简单的类对象，包含三个属性：`flags`、`type` 和 `target`。
CustomArgs = collections.namedtuple('CustomArgs', 'flags type target')
options = [
    CustomArgs(['--lr', '--learning_rate'], type=float, target=('optimizer', 'args', 'lr')),
    CustomArgs(['--bs', '--batch_size'], type=int, target=('data_loader', 'args', 'batch_size'))
    # 在此处添加的选项可以通过命令行标志进行修改。
]

target 参数应为一个键序列，用于在配置字典中访问该选项。例如，在本例中，学习率选项的 target 是 ('optimizer', 'args', 'lr')，因为 config['optimizer']['args']['lr'] 指向学习率。运行 python train.py -c config.json --bs 256 将使用 config.json 中的配置进行训练，但批量大小会被命令行选项覆盖，调整为 256。

数据加载器

• 编写自定义数据加载器

1. 继承 BaseDataLoader

   BaseDataLoader 是 torch.utils.data.DataLoader 的子类，你可以选择使用其中任一。
   BaseDataLoader 负责：

   • 生成下一个批次
   • 数据打乱顺序
   • 通过调用 BaseDataLoader.split_validation() 生成验证数据加载器

• 数据加载器的使用

  BaseDataLoader 是一个迭代器，遍历批次的方式如下：

  for batch_idx, (x_batch, y_batch) in data_loader:
      pass
  

• 示例

  请参考 data_loader/data_loaders.py 中的 MNIST 数据加载示例。

训练器

• 编写自定义训练器

1. 继承 BaseTrainer

   BaseTrainer 负责：

   • 训练过程的日志记录
   • 模型检查点的保存
   • 检查点的恢复
   • 可配置的性能监控，用于保存当前最佳模型并提前停止训练。
     • 如果配置中的 monitor 设置为 max val_accuracy，则当某个 epoch 的验证准确率超过当前最高值时，训练器会保存一个名为 model_best.pth 的检查点。
     • 如果配置了 early_stop，当模型性能在指定的 epoch 数内没有提升时，训练将自动终止。可通过将 early_stop 选项设为 0 或直接删除该配置行来关闭此功能。

2. 实现抽象方法

   你需要实现 _train_epoch() 方法以完成你的训练流程；如果需要验证，还可以像 trainer/trainer.py 中那样实现 _valid_epoch() 方法。

• 示例

  请参考 trainer/trainer.py 中的 MNIST 训练示例。

• 基于迭代的训练

  Trainer.__init__ 接受一个可选参数 len_epoch，用于控制每个 epoch 中的批次（步骤）数量。

模型

• 编写自定义模型

1. 继承 BaseModel

   BaseModel 负责：

   • 继承自 torch.nn.Module
   • __str__：修改原生 print 函数，使其输出可训练参数的数量。

2. 实现抽象方法

   实现前向传播方法 forward()。

• 示例

  请参考 model/model.py 中的 LeNet 示例。

损失函数

自定义损失函数可以在 model/loss.py 中实现。通过将配置文件中 loss 字段的名称更改为相应名称即可使用。

评估指标

评估指标函数位于 model/metric.py 中。

你可以在配置文件中提供一个指标列表来监控多个指标，例如：

"metrics": ["accuracy", "top_k_acc"],

额外日志记录

如果你有需要记录的额外信息，可在训练器类的 _train_epoch() 方法中，将其与 log 合并后再返回，如下所示：

additional_log = {"gradient_norm": g, "sensitivity": s}
log.update(additional_log)
return log

测试

你可以通过运行 test.py 并使用 --resume 参数传入已训练好的检查点路径来测试模型。

验证数据

要从数据加载器中分离出验证数据，可以调用 BaseDataLoader.split_validation()，它将返回一个大小由配置文件指定的验证数据加载器。 validation_split 可以是验证集占总数据的比例（0.0 ≤ float < 1.0），也可以是样本数量（0 ≤ int < n_total_samples）。

注意：split_validation() 方法会修改原始数据加载器。
注意：如果 "validation_split" 设置为 0，split_validation() 将返回 None。

检查点

你可以在配置文件中指定训练会话的名称：

"name": "MNIST_LeNet",

检查点将保存在 save_dir/name/timestamp/checkpoint_epoch_n 目录下，时间戳格式为 mmdd_HHMMSS。 同一目录下还会保存一份配置文件副本。

注意：检查点包含以下内容：

{
  'arch': arch,
  'epoch': epoch,
  'state_dict': self.model.state_dict(),
  'optimizer': self.optimizer.state_dict(),
  'monitor_best': self.mnt_best,
  'config': self.config
}

TensorBoard 可视化

此模板支持使用 torch.utils.tensorboard 或 TensorboardX 进行 TensorBoard 可视化。

1. 安装

   如果你使用的是 PyTorch 1.1 或更高版本，可以通过 pip install tensorboard>=1.14.0 安装 TensorBoard。

   否则，建议安装 TensorboardX。请按照 TensorboardX 的安装指南进行操作。

2. 运行训练

   确保配置文件中的 tensorboard 选项已开启。

   "tensorboard" : true
   

3. 启动 TensorBoard 服务器

   在项目根目录下输入 tensorboard --logdir saved/log/，服务器将在 http://localhost:6006 打开。

   默认情况下，配置文件中指定的损失和指标、输入图像以及模型参数直方图都会被记录。若需更多可视化内容，可在 trainer._train_epoch 方法中使用 add_scalar('tag', data)、add_image('tag', image) 等方法。 此模板中的 add_something() 方法实际上是 tensorboardX.SummaryWriter 和 torch.utils.tensorboard.SummaryWriter 模块中相应方法的封装。

   注意：你无需手动指定当前步数，因为 logger/visualization.py 中定义的 WriterTensorboard 类会自动跟踪当前步数。

贡献

欢迎贡献任何功能或改进！代码风格遵循 PEP8 规范。 提交代码前，请确保通过 Flake8 检查。

待办事项

• 多个优化器
• 支持更多 TensorBoard 功能
• 使用固定的随机种子
• 支持 PyTorch 原生 TensorBoard
• tensorboardX 日志支持
• 可配置的日志布局和检查点命名
• 基于迭代的训练（而非基于轮次的训练）
• 添加用于微调的命令行选项

许可证

本项目采用 MIT 许可证。详情请参阅 LICENSE 文件。

致谢

本项目受到 Mahmoud Gemy 的项目 Tensorflow-Project-Template 的启发。

PyTorch Template 快速上手指南

环境准备

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

• 操作系统: Linux, macOS 或 Windows
• Python: >= 3.5 (推荐 3.6+)
• PyTorch: >= 0.4 (推荐 1.2+)
• 可选依赖:
  • tqdm: 用于显示训练进度条
  • tensorboard >= 1.14: 用于可视化训练过程（若使用 PyTorch 1.1+）
  • tensorboardX: 若使用较低版本的 PyTorch，需安装此库

国内加速建议： 安装依赖时，推荐使用清华或阿里镜像源以提升下载速度：

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

本项目是一个模板工程，无需通过 pip 安装，直接克隆代码即可使用。

1. 克隆仓库

   git clone https://github.com/victoresque/pytorch-template.git
   cd pytorch-template
   

2. 安装依赖 项目根目录下通常包含 requirements.txt，执行以下命令安装所需包：

   pip install -r requirements.txt
   

   若没有 requirements.txt，请手动安装核心依赖：

   pip install torch torchvision tqdm tensorboard
   

3. 初始化新项目（可选） 如果您想基于此模板创建一个全新的独立项目文件夹（自动清理无关文件），可运行：

   python new_project.py ../MyNewProject
   cd ../MyNewProject
   

基本使用

本模板默认包含一个 MNIST 分类示例，通过修改 JSON 配置文件即可轻松调整超参数和训练流程。

1. 配置文件说明

所有训练参数均在 config.json 中定义，包括模型架构、数据加载器、优化器、学习率策略及训练轮数等。您可以根据需要编辑该文件。

关键配置项示例：

{
  "name": "Mnist_LeNet",
  "n_gpu": 1,
  "arch": {
    "type": "MnistModel"
  },
  "data_loader": {
    "type": "MnistDataLoader",
    "args": {
      "batch_size": 64,
      "validation_split": 0.1
    }
  },
  "trainer": {
    "epochs": 100,
    "save_dir": "saved/",
    "tensorboard": true
  }
}

2. 开始训练

使用默认配置文件启动训练：

python train.py -c config.json

常用命令行选项：

• 指定配置文件：

  python train.py --config config.json
  

• 覆盖配置参数（无需修改 json 文件即可快速调整）：

  # 例如：将批大小改为 256，学习率改为 0.0001
  python train.py -c config.json --bs 256 --lr 0.0001
  

• 多 GPU 训练： 在 config.json 中设置 "n_gpu" 大于 1，或通过环境变量指定显卡：

  CUDA_VISIBLE_DEVICES=0,1 python train.py -c config.json
  

  或在命令中直接指定：

  python train.py --device 0,1 -c config.json
  

3. 断点续训

如果训练意外中断，可以通过指定检查点路径恢复训练：

python train.py --resume saved/models/Mnist_LeNet/0520_143022/checkpoint_epoch_10.pth

4. 模型测试

训练完成后，使用 test.py 评估模型性能：

python test.py --resume saved/models/Mnist_LeNet/0520_143022/checkpoint_epoch_100.pth

5. 可视化监控

若配置文件中开启了 "tensorboard": true，训练日志将自动保存。启动 TensorBoard 查看损失曲线和准确率：

tensorboard --logdir saved/log

然后在浏览器访问 http://localhost:6006 即可查看可视化图表。