原文链接:Tutorial 2: Finetuning Models — mmflow documentation
教程2:微调模型
在FlyingChairs和FlyingThings3d数据集上预训练的光流估计器可以作为其他数据集的较好的预训练模型。本教程为用户提供了如何利用”模型园“Model Zoohttps://mmflow.readthedocs.io/en/latest/model_zoo.html中模型在其他数据集上训练,从而获得更好表现的指导说明。
修改训练策略
微调超参数和默认的训练策略是不同的,通常需要更小的学习率和更少的训练迭代次数。
# 优化器
optimizer = dict(type='Adam', lr=1e-5, weight_decay=0.0004, betas=(0.9, 0.999))
optimizer_config = dict(grad_clip=None)
# learning policy
lr_config = dict(
policy='step',
by_epoch=False,
gamma=0.5,
step=[
45000, 65000, 85000, 95000, 97500, 100000, 110000, 120000, 130000,
140000
])
runner = dict(type='IterBasedRunner', max_iters=150000)
checkpoint_config = dict(by_epoch=False, interval=10000)
evaluation = dict(interval=10000, metric='EPE')
使用预训练模型
用户可以加载一个预训练模型,通过设定load_from字段,将其指向模型的路径即可。用户需要在训练前下载好模型权重参数文件,避免在训练过程中下载。
# 使用PWC-Net的预训练模型
load_from = 'https://download.openmmlab.com/mmflow/pwcnet/pwcnet_8x1_sfine_flyingthings3d_subset_384x768.pth'
在单个GPU上训练
我们提供了tools/train.py来在单个GPU上发起训练任务。基础用法如下:
python tools/train.py ${CONFIG_FILE} [optional arguments]
在训练过程中,日志文件和检查点文件将被保存在工作目录,该目录由配置文件中的work_dir指定或者通过CLI参数--work-dir指定。
这一工具支持多个可选参数,包括:
-
--work-dir ${WORK_DIR}
: 覆盖工作目录. -
--resume-from ${CHECKPOINT_FILE}
: 从之前的检查点文件中恢复. -
--cfg-option
: 覆写配置文件中的一些设置, xxx=yyy格式的键值对将被写入配置文件中.例如: ‘–cfg-option model.encoder.in_channels=6’.
注意:
resume-from和load-from的区别:
resume-from加载了模型权重参数和优化器的状态,迭代过程也从指定的检查点文件中继承。通常被用于从偶然中断的训练过程中进行恢复。而load-from加载了模型权重参数并且从0开始迭代训练,通常用于微调。
在CPU上训练
在CPU上训练的过程和在单GPU上训练的过程一致,我们只需要在训练前禁用GPU,
export CUDA_VISIBLE_DEVICES=-1
然后运行上述脚本即可。
我们不推荐使用CPU训练,因为太慢了!我们支持这个功能只是为了方便用户在没有GPU的电脑上进行调试。
在多GPU上训练
MMFlow基于MMDistributedDataParallel实现了分布式训练。
我们提供了tools/dist_train.sh来发起在多个GPU上的训练,基础用法如下:
sh tools/dist_train.sh ${CONFIG_FILE} ${GPU_NUM} [optional arguments]
优化器参数和上述的单GPU训练的相同,并且有额外的参数用于指定GPU数量。
在单台机器上发起多个任务
如果你项在单台机器上发起多个任务,例如在一台有8个GPU的设备上利用4个GPU训练两个任务(这是什么人家啊,土豪吧=_=),你需要为每个任务指定不同的端口(默认29500),从而避免通信访问冲突。
如果你用dist_train.sh来发起任务,你可以用命令行指定接口。
CUDA_VISIBLE_DEVICES=0,1,2,3 PORT=29500 sh tools/dist_train.sh ${CONFIG_FILE} 4
CUDA_VISIBLE_DEVICES=4,5,6,7 PORT=29501 sh tools/dist_train.sh ${CONFIG_FILE} 4
在多节点上训练
MMFlow依赖于torch.distributed包实现分布式训练,因此,一个基础的用法是,通过Pytorch的启动工具Distributed communication package - torch.distributed — PyTorch 1.13 documentation来发起分布式训练
在多台设备上训练
如果你在多台通过以色网连接的设备上发起任务,你可以通过下面的命令简单的运行:
在第一台设备上:
NNODES=2 NODE_RANK=0 PORT=${MASTER_PORT} MASTER_ADDR=${MASTER_ADDR} sh tools/dist_train.sh ${CONFIG_FILE} ${GPUS}
在第二台设备上:
NNODES=2 NODE_RANK=1 PORT=${MASTER_PORT} MASTER_ADDR=${MASTER_ADDR} sh tools/dist_train.sh ${CONFIG_FILE} ${GPUS}
通常情况下如果没有像IfiniBand之类的高速网络连接的话,速度是比较慢的。
使用Slurm管理任务
Slurm Workload Manager - DocumentationSlurm是一个计算集群的优秀的任务调度系统。在Slurm管理的计算集群上,你可以通过slurm_train.sh来分发任务。(这段话好眼熟,没错就是上一篇博客=_=) 它支持单节点和多节点训练。
基础用法如下:
[GPUS=${GPUS}] sh tools/slurm_train.sh ${PARTITION} ${JOB_NAME} ${CONFIG_FILE} ${WORK_DIR}
下面是一个用户名为dev,在8GPU上训练PWC-Net,工作目录指向相同的文件系统目录的使用例子:
GPUS=8 sh tools/slurm_train.sh dev pwc_chairs configs/pwcnet/pwcnet_8x1_slong_flyingchairs_384x448.py work_dir/pwc_chairs
你可以通过查看源码来了解全部的参数和环境变量。当使用Slurm时,端口需要给定,可以通过下列方式:
1. 通过--cfg-options设定端口。推荐使用该方法,因为它不会改变原始配置。
GPUS=4 GPUS_PER_NODE=4 sh tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config1.py ${WORK_DIR} --cfg-options 'dist_params.port=29500'
GPUS=4 GPUS_PER_NODE=4 sh tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config2.py ${WORK_DIR} --cfg-options 'dist_params.port=29501'
2.修改配置文件,从而指定不同的通信端口
在config1.py中
dist_params = dict(backend='nccl', port=29500)
在config2.py中
dist_params = dict(backend='nccl', port=29501)
然后你可以利用config1.py和config2.py发起任务:
GPUS=4 GPUS_PER_NODE=4 sh tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config1.py ${WORK_DIR}
GPUS=4 GPUS_PER_NODE=4 sh tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config2.py ${WORK_DIR}