Spaces:
Sleeping
Sleeping
File size: 18,273 Bytes
14c9181 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 |
# 训练与测试
为了适配多样化的用户需求，MMOCR 实现了多种不同操作系统及设备上的模型训练及测试。无论是使用本地机器进行单机单卡训练测试，还是在部署了 slurm 系统的大规模集群上进行训练测试，MMOCR 都提供了便捷的解决方案。
## 单卡机器训练及测试
### 训练
`tools/train.py` 实现了基础的训练服务。MMOCR 推荐用户使用 GPU 进行模型训练和测试，但是，用户也可以通过指定 `CUDA_VISIBLE_DEVICES=-1` 来使用 CPU 设备进行模型训练及测试。例如，以下命令演示了如何使用 CPU 或单卡 GPU 来训练 DBNet 文本检测器。
```bash
# 通过调用 tools/train.py 来训练指定的 MMOCR 模型
CUDA_VISIBLE_DEVICES= python tools/train.py ${CONFIG_FILE} [PY_ARGS]
# 训练
# 示例 1：使用 CPU 训练 DBNet
CUDA_VISIBLE_DEVICES=-1 python tools/train.py configs/textdet/dbnet/dbnet_resnet50-dcnv2_fpnc_1200e_icdar2015.py
# 示例 2：指定使用 gpu:0 训练 DBNet，指定工作目录为 dbnet/，并打开混合精度（amp）训练
CUDA_VISIBLE_DEVICES=0 python tools/train.py configs/textdet/dbnet/dbnet_resnet50-dcnv2_fpnc_1200e_icdar2015.py --work-dir dbnet/ --amp
```
```{note}
此外，如需使用指定编号的 GPU 进行训练或测试，例如使用3号 GPU，则可以通过设定 CUDA_VISIBLE_DEVICES=3 来实现。
```
下表列出了 `train.py` 支持的所有参数。其中，不带 `--` 前缀的参数为必须的位置参数，带 `--` 前缀的参数为可选参数。
| 参数 | 类型 | 说明 |
| --------------- | ---- | -------------------------------------------------------------- |
| config | str | （必须）配置文件路径。 |
| --work-dir | str | 指定工作目录，用于存放训练日志以及模型 checkpoints。 |
| --resume | bool | 是否从断点处恢复训练。 |
| --amp | bool | 是否使用混合精度。 |
| --auto-scale-lr | bool | 是否使用学习率自动缩放。 |
| --cfg-options | str | 用于覆写配置文件中的指定参数。[示例](#添加示例) |
| --launcher | str | 启动器选项，可选项目为 \['none', 'pytorch', 'slurm', 'mpi'\]。 |
| --local_rank | int | 本地机器编号，用于多机多卡分布式训练，默认为 0。 |
### 测试
`tools/test.py` 提供了基础的测试服务，其使用原理和训练脚本类似。例如，以下命令演示了 CPU 或 GPU 单卡测试 DBNet 模型。
```bash
# 通过调用 tools/test.py 来测试指定的 MMOCR 模型
CUDA_VISIBLE_DEVICES= python tools/test.py ${CONFIG_FILE} ${CHECKPOINT_FILE} [PY_ARGS]
# 测试
# 示例 1：使用 CPU 测试 DBNet
CUDA_VISIBLE_DEVICES=-1 python tools/test.py configs/textdet/dbnet/dbnet_resnet50-dcnv2_fpnc_1200e_icdar2015.py dbnet_r50.pth
# 示例 2：使用 gpu:0 测试 DBNet
CUDA_VISIBLE_DEVICES=0 python tools/test.py configs/textdet/dbnet/dbnet_resnet50-dcnv2_fpnc_1200e_icdar2015.py dbnet_r50.pth
```
下表列出了 `test.py` 支持的所有参数。其中，不带 `--` 前缀的参数为必须的位置参数，带 `--` 前缀的参数为可选参数。
| 参数 | 类型 | 说明 |
| ------------- | ----- | -------------------------------------------------------------- |
| config | str | （必须）配置文件路径。 |
| checkpoint | str | （必须）待测试模型路径。 |
| --work-dir | str | 工作目录，用于存放训练日志以及模型 checkpoints。 |
| --save-preds | bool | 是否将预测结果写入 pkl 文件并保存。 |
| --show | bool | 是否可视化预测结果。 |
| --show-dir | str | 将可视化的预测结果保存至指定路径。 |
| --wait-time | float | 可视化间隔时间（秒），默认为 2 秒。 |
| --cfg-options | str | 用于覆写配置文件中的指定参数。[示例](#添加示例) |
| --launcher | str | 启动器选项，可选项目为 \['none', 'pytorch', 'slurm', 'mpi'\]。 |
| --local_rank | int | 本地机器编号，用于多机多卡分布式训练，默认为 0。 |
| --tta | bool | 是否使用测试时数据增强 |
## 多卡机器训练及测试
对于大规模模型，采用多 GPU 训练和测试可以极大地提升操作的效率。为此，MMOCR 提供了基于 [MMDistributedDataParallel](mmengine.model.wrappers.MMDistributedDataParallel) 实现的分布式脚本 `tools/dist_train.sh` 和 `tools/dist_test.sh`。
```bash
# 训练
NNODES=${NNODES} NODE_RANK=${NODE_RANK} PORT=${MASTER_PORT} MASTER_ADDR=${MASTER_ADDR} ./tools/dist_train.sh ${CONFIG_FILE} ${GPU_NUM} [PY_ARGS]
# 测试
NNODES=${NNODES} NODE_RANK=${NODE_RANK} PORT=${MASTER_PORT} MASTER_ADDR=${MASTER_ADDR} ./tools/dist_test.sh ${CONFIG_FILE} ${CHECKPOINT_FILE} ${GPU_NUM} [PY_ARGS]
```
下表列出了 `dist_*.sh` 支持的参数：
| 参数 | 类型 | 说明 |
| --------------- | ---- | ---------------------------------------------------------------------------------- |
| NNODES | int | 总共使用的机器节点个数，默认为 1。 |
| NODE_RANK | int | 节点编号，默认为 0。 |
| PORT | int | 在 RANK 0 机器上使用的 MASTER_PORT 端口号，取值范围是 0 至 65535，默认值为 29500。 |
| MASTER_ADDR | str | RANK 0 机器的 IP 地址，默认值为 127.0.0.1。 |
| CONFIG_FILE | str | （必须）指定配置文件的地址。 |
| CHECKPOINT_FILE | str | （必须，仅在 dist_test.sh 中适用）指定模型权重的地址。 |
| GPU_NUM | int | （必须）指定 GPU 的数量。 |
| \[PY_ARGS\] | str | 该部分一切的参数都会被直接传入 tools/train.py 或 tools/test.py 中。 |
这两个脚本可以实现**单机多卡**或**多机多卡**的训练和测试，下面演示了它们在不同场景下的用法。
### 单机多卡
以下命令演示了如何在搭载多块 GPU 的**单台机器**上使用指定数目的 GPU 进行训练及测试：
1. **训练**
使用单台机器上的 4 块 GPU 训练 DBNet。
```bash
# 单机 4 卡训练 DBNet
tools/dist_train.sh configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py 4
```
2. **测试**
使用单台机器上的 4 块 GPU 测试 DBNet。
```bash
# 单机 4 卡测试 DBNet
tools/dist_test.sh configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py dbnet_r50.pth 4
```
### 单机多任务训练及测试
对于搭载多块 GPU 的单台服务器而言，用户可以通过指定 GPU 的形式来同时执行不同的训练任务。例如，以下命令演示了如何在一台 8 卡 GPU 服务器上分别使用 `[0, 1, 2, 3]` 卡测试 DBNet 及 `[4, 5, 6, 7]` 卡训练 CRNN：
```bash
# 指定使用 gpu:0,1,2,3 测试 DBNet，并分配端口号 29500
CUDA_VISIBLE_DEVICES=0,1,2,3 PORT=29500 ./tools/dist_test.sh configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py dbnet_r50.pth 4
# 指定使用 gpu:4,5,6,7 训练 CRNN，并分配端口号 29501
CUDA_VISIBLE_DEVICES=4,5,6,7 PORT=29501 ./tools/dist_train.sh configs/textrecog/crnn/crnn_academic_dataset.py 4
```
```{note}
`dist_train.sh` 默认将 `MASTER_PORT` 设置为 `29500`，当单台机器上有其它进程已占用该端口时，程序则会出现运行时错误 `RuntimeError: Address already in use`。此时，用户需要将 `MASTER_PORT` 设置为 `(0~65535)` 范围内的其它空闲端口号。
```
### 多机多卡训练及测试
MMOCR 基于[torch.distributed](https://pytorch.org/docs/stable/distributed.html#launch-utility) 提供了相同局域网下的多台机器间的多卡分布式训练。
1. **训练**
以下命令演示了如何在两台机器上分别使用 2 张 GPU 合计 4 卡训练 DBNet：
```bash
# 示例：在两台机器上分别使用 2 张 GPU 合计 4 卡训练 DBNet
# 在 “机器1” 上运行以下命令
NNODES=2 NODE_RANK=0 PORT=29501 MASTER_ADDR=10.140.0.169 tools/dist_train.sh configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py 2
# 在 “机器2” 上运行以下命令
NNODES=2 NODE_RANK=1 PORT=29501 MASTER_ADDR=10.140.0.169 tools/dist_train.sh configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py 2
```
2. **测试**
以下命令演示了如何在两台机器上分别使用 2 张 GPU 合计 4 卡测试：
```bash
# 示例：在两台机器上分别使用 2 张 GPU 合计 4 卡测试
# 在 “机器1” 上运行以下命令
NNODES=2 NODE_RANK=0 PORT=29500 MASTER_ADDR=10.140.0.169 tools/dist_test.sh configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py dbnet_r50.pth 2
# 在 “机器2” 上运行以下命令
NNODES=2 NODE_RANK=1 PORT=29501 MASTER_ADDR=10.140.0.169 tools/dist_test.sh configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py dbnet_r50.pth 2
```
```{note}
需要注意的是，采用多机多卡训练时，机器间的网络传输速度可能成为训练速度的瓶颈。
```
## 集群训练及测试
针对 [Slurm](https://slurm.schedmd.com/) 调度系统管理的计算集群，MMOCR 提供了对应的训练和测试任务提交脚本 `tools/slurm_train.sh` 及 `tools/slurm_test.sh`。
```bash
# tools/slurm_train.sh 提供基于 slurm 调度系统管理的计算集群上提交训练任务的脚本
GPUS=${GPUS} GPUS_PER_NODE=${GPUS_PER_NODE} CPUS_PER_TASK=${CPUS_PER_TASK} SRUN_ARGS=${SRUN_ARGS} ./tools/slurm_train.sh ${PARTITION} ${JOB_NAME} ${CONFIG_FILE} ${WORK_DIR} [PY_ARGS]
# tools/slurm_test.sh 提供基于 slurm 调度系统管理的计算集群上提交测试任务的脚本
GPUS=${GPUS} GPUS_PER_NODE=${GPUS_PER_NODE} CPUS_PER_TASK=${CPUS_PER_TASK} SRUN_ARGS=${SRUN_ARGS} ./tools/slurm_test.sh ${PARTITION} ${JOB_NAME} ${CONFIG_FILE} ${CHECKPOINT_FILE} ${WORK_DIR} [PY_ARGS]
```
| 参数 | 类型 | 说明 |
| --------------- | ---- | ------------------------------------------------------------------------- |
| GPUS | int | 使用的 GPU 数目，默认为8。 |
| GPUS_PER_NODE | int | 每台节点机器上搭载的 GPU 数目，默认为8。 |
| CPUS_PER_TASK | int | 任务使用的 CPU 个数，默认为5。 |
| SRUN_ARGS | str | 其他 srun 支持的参数。详见[这里](https://slurm.schedmd.com/srun.html) |
| PARTITION | str | （必须）指定使用的集群分区。 |
| JOB_NAME | str | （必须）提交任务的名称。 |
| WORK_DIR | str | （必须）任务的工作目录，训练日志以及模型的 checkpoints 将被保存至该目录。 |
| CHECKPOINT_FILE | str | （必须，仅在 slurm_test.sh 中适用）指向模型权重的地址。 |
| \[PY_ARGS\] | str | tools/train.py 以及 tools/test.py 支持的参数。 |
这两个脚本可以实现 slurm 集群上的训练和测试，下面演示了它们在不同场景下的用法。
1. 训练
以下示例为在 slurm 集群 dev 分区申请 1 块 GPU 进行 DBNet 训练。
```bash
# 示例：在 slurm 集群 dev 分区申请 1块 GPU 资源进行 DBNet 训练任务
GPUS=1 GPUS_PER_NODE=1 CPUS_PER_TASK=5 tools/slurm_train.sh dev db_r50 configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py work_dir
```
2. 测试
同理， 则提供了测试任务提交脚本。以下示例为在 slurm 集群 dev 分区申请 1 块 GPU 资源进行 DBNet 测试。
```bash
# 示例：在 slurm 集群 dev 分区申请 1块 GPU 资源进行 DBNet 测试任务
GPUS=1 GPUS_PER_NODE=1 CPUS_PER_TASK=5 tools/slurm_test.sh dev db_r50 configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py dbnet_r50.pth work_dir
```
## 进阶技巧
### 从断点恢复训练
`tools/train.py` 提供了从断点恢复训练的功能，用户仅需在命令中指定 `--resume` 参数，即可自动从断点恢复训练。
```bash
# 示例：从断点恢复训练
python tools/train.py configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py 4 --resume
```
默认地，程序将自动从上次训练过程中最后成功保存的断点，即 `latest.pth` 处开始继续训练。如果用户希望指定从特定的断点处开始恢复训练，则可以按如下格式在模型的配置文件中设定该断点的路径。
```python
# 示例：在配置文件中设置想要加载的断点路径
load_from = 'work_dir/dbnet/models/epoch_10000.pth'
```
### 混合精度训练
混合精度训练可以在缩减内存占用的同时提升训练速度，为此，MMOCR 提供了一键式的混合精度训练方案，仅需在训练时添加 `--amp` 参数即可。
```bash
# 示例：使用自动混合精度训练
python tools/train.py configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py 4 --amp
```
下表列出了 MMOCR 中各算法对自动混合精度训练的支持情况：
| | 是否支持混合精度训练 | 备注 |
| ------------- | :------------------: | :---------------------------: |
| | 文本检测 | |
| DBNet | 是 | |
| DBNetpp | 是 | |
| DRRG | 否 | roi_align_rotated 不支持 fp16 |
| FCENet | 否 | BCELoss 不支持 fp16 |
| Mask R-CNN | 是 | |
| PANet | 是 | |
| PSENet | 是 | |
| TextSnake | 否 | |
| | 文本识别 | |
| ABINet | 是 | |
| ASTER | 是 | |
| CRNN | 是 | |
| MASTER | 是 | |
| NRTR | 是 | |
| RobustScanner | 是 | |
| SAR | 是 | |
| SATRN | 是 | |
### 自动学习率缩放
MMOCR 在配置文件中为每一个模型设置了默认的初始学习率，然而，当用户使用的 `batch_size` 不同于我们预设的 `base_batch_size` 时，这些初始学习率可能不再完全适用。因此，我们提供了自动学习率缩放工具。当使用不同于 MMOCR 预设的 `base_batch_size` 进行训练时，用户仅需添加 `--auto-scale-lr` 参数即可自动依据新的 `batch_size` 将学习率缩放至对应尺度。
```bash
# 示例：使用自动学习率缩放
python tools/train.py configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py 4 --auto-scale-lr
```
### 可视化模型测试结果
`tools/test.py` 提供了可视化接口，以方便用户对模型进行定性分析。
<div align="center">
(绿色框为真实标注，红色框为预测结果)
</div>
<div align="center">
（绿色字体为真实标注，红色字体为预测结果）
</div>
<div align="center">
（从左至右分别为：原图，文本检测和识别结果，文本分类结果，关系图）
</div>
```bash
# 示例1：每间隔 2 秒绘制出
python tools/test.py configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py dbnet_r50.pth --show --wait-time 2
# 示例2：对于不支持图形化界面的系统（如计算集群等），可以将可视化结果存入指定路径
python tools/test.py configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py dbnet_r50.pth --show-dir ./vis_results
```
`tools/test.py` 中可视化相关参数说明：
| 参数 | 类型 | 说明 |
| ----------- | ----- | -------------------------------- |
| --show | bool | 是否绘制可视化结果。 |
| --show-dir | str | 可视化图片存储路径。 |
| --wait-time | float | 可视化间隔时间（秒），默认为 2。 |
### 测试时数据增强
测试时增强，指的是在推理（预测）阶段，将原始图片进行水平翻转、垂直翻转、对角线翻转、旋转角度等数据增强操作，得到多张图，分别进行推理，再对多个结果进行综合分析，得到最终输出结果。
为此，MMOCR 提供了一键式测试时数据增强，仅需在测试时添加 `--tta` 参数即可。
```{note}
TTA 仅支持文本识别模型。
```
```bash
python tools/test.py configs/textrecog/crnn/crnn_mini-vgg_5e_mj.py checkpoints/crnn_mini-vgg_5e_mj.pth --tta
```
|