教程 3: 使用现有模型进行预训练¶
本文档提供有关如何运行算法以及如何使用 MMSelfSup 中的一些工具的基本用法。有关安装说明和数据准备，请参阅 install.md 和 prepare_data.md。
## 开始训练
注意：配置文件中的默认学习率是针对特定数量的 GPU（GPU数量已在配置文件名称中注明）。如果使用不同数量的 GPUs，总的 batch size 将按比例变化，您必须按照 new_lr = old_lr * new_ngpus / old_ngpus 缩放学习率。
使用单卡训练¶
python tools/train.py ${CONFIG_FILE} [optional arguments]
一个简单的例子来开启训练：
python tools/train.py configs/selfsup/mae/mae_vit-base-p16_8xb512-coslr-400e_in1k.py
使用 CPU 训练¶
export CUDA_VISIBLE_DEVICES=-1
python tools/train.py ${CONFIG_FILE} [optional arguments]
注意。我们不建议用户使用CPU进行训练，因为它太慢了。我们支持这个功能，是为了方便用户在没有GPU的机器上进行调试。
使用多卡训练¶
sh tools/dist_train.sh ${CONFIG_FILE} ${GPUS} [optional arguments]
可选参数：
--work-dir：指示您的自定义工作目录以保存 checkpoints 和日志。--resume：自动在你的工作目录中查找最新的 checkpoints。或者设置--resume ${CHECKPOINT_PATH}来加载特定的 checkpoints 文件。--amp：启用自动混合精度训练。--cfg-options：设置--cfg-options将修改原始配置。例如，设置--cfg-options randomness.seed=0将为随机数设置种子。
使用 8 个 GPUs 开始训练的示例：
sh tools/dist_train.sh configs/selfsup/mae/mae_vit-base-p16_8xb512-coslr-400e_in1k.py 8
或者，如果您在使用 slurm 管理的集群上运行 MMSelfSup：
GPUS_PER_NODE=${GPUS_PER_NODE} GPUS=${GPUS} SRUN_ARGS=${SRUN_ARGS} sh tools/slurm_train.sh ${PARTITION} ${JOB_NAME} ${CONFIG_FILE} [optional arguments]
使用 8 个 GPUs 开始训练的示例：
# 默认设置: GPUS_PER_NODE=8 GPUS=8
sh tools/slurm_train.sh Dummy Test_job configs/selfsup/mae/mae_vit-base-p16_8xb512-coslr-400e_in1k.py
使用用多台机器训练¶
如果您想使用由以太网连接起来的多台机器， 您可以使用以下命令:
在第一台机器上:
NNODES=2 NODE_RANK=0 PORT=${MASTER_PORT} MASTER_ADDR=${MASTER_ADDR} sh tools/dist_train.sh ${CONFIG} ${GPUS}
在第二台机器上:
NNODES=2 NODE_RANK=1 PORT=${MASTER_PORT} MASTER_ADDR=${MASTER_ADDR} sh tools/dist_train.sh ${CONFIG} ${GPUS}
但是，如果您不使用高速网路连接这几台机器的话，训练将会非常慢。
如果您使用的是 slurm 来管理多台机器，您可以使用同在单台机器上一样的命令来启动任务，但是您必须得设置合适的环境变量和参数，具体可以参考 slurm_train.sh。
在一台机器上启动多个任务¶
如果你在一台机器上启动多个任务，例如，在一台有 8 个 GPU 的机器上启动 2 个分别使用４块 GPU 的训练任务，你需要为每个任务指定不同的端口（默认为29500）以避免通信冲突。
如果你使用 dist_train.sh 来启动训练任务。
CUDA_VISIBLE_DEVICES=0,1,2,3 PORT=29500 sh tools/dist_train.sh ${CONFIG_FILE} 4 --work-dir tmp_work_dir_1
CUDA_VISIBLE_DEVICES=4,5,6,7 PORT=29501 sh tools/dist_train.sh ${CONFIG_FILE} 4 --work-dir tmp_work_dir_2
如果你用 slurm 启动训练任务，你有两种方法来设置不同的通信端口。
方法 1 :
在 config1.py 中:
env_cfg = dict(dist_cfg=dict(backend='nccl', port=29500))
在 config2.py 中:
env_cfg = dict(dist_cfg=dict(backend='nccl', port=29501))
然后你可以用 config1.py 和 config2.py 启动两个任务。
CUDA_VISIBLE_DEVICES=0,1,2,3 GPUS=4 sh tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config1.py [optional arguments]
CUDA_VISIBLE_DEVICES=4,5,6,7 GPUS=4 sh tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config2.py [optional arguments]
方法 2 :
你可以设置不同的通信端口，而不需要修改配置文件，但必须设置 --cfg-options 来覆盖配置文件中的默认端口。
CUDA_VISIBLE_DEVICES=0,1,2,3 GPUS=4 sh tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config1.py --work-dir tmp_work_dir_1 --cfg-options env_cfg.dist_cfg.port=29500
CUDA_VISIBLE_DEVICES=4,5,6,7 GPUS=4 sh tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config2.py --work-dir tmp_work_dir_2 --cfg-options env_cfg.dist_cfg.port=29501