本指导适用于在 TencentOS Server 3上使用 Distributed DataParallel(DDP)训练框架运行 ViT(Vision Transformer)模型,以 Docker 方式启动。
前置环境条件
运行模型
DDP 运行的文件为 distributed.py,其中有以下几处需要修改。
代码第74行,将 --local_rank 参数的 default 改为 int(os.environ["LOCAL_RANK"],以便启用最新的 torchrun 方法启动。
#原有代码default=-1,#更改为default=int(os.environ["LOCAL_RANK"]),
代码第375行会报错,将 view() 方法更改为 reshape() 方法即可。
#原有代码correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)#更改为correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
以上修改完之后,即可运行模型。这里我们测试的是 vit_large_16 模型通过 ImageNet-1K 数据集训练。由于我们有8张L40 gpu,所以 --nproc_per_node 设置为8。
torchrun --nproc_per_node=8 distributed.py --data datasets/imagenet --arch vit_l_16 --epochs 400 --batch-size 768 --learning-rate 0.001
此时会使用 vit_large_16 模型架构,从0开始训练400个 epochs,数据集在 datasets/imagenet 文件夹下,batchsize 为768,学习率 lr 为0.001,优化方法为SGD 使用 momentum=0.9,权重衰退 weight decay=1e-4。部分参数设置来源于 ViT 原论文:ViT,从而尽可能的复现出论文里的结果。
说明:
如需要更改 momentum 参数和权重衰减参数,只需加上对应参数
--momentum
,--weight-decay
并设置想要的值即可。同时测试环境L40如果 batchsize 选择1024则会显存溢出,所以选择768。
此时会开始训练模型,由于使用的是 DDP 训练模式,所以各个 gpu 之间的显存占用相差不大。运行模型过程中会保存当前 epoch 训练结束之后的模型权重以及到当前 epoch 为止在验证集上最好性能的 epoch 的模型权重。
训练时命令行会出现训练时每一个 teration 和测试时每一个 iteration 的性能以及每一个 epoch 测试完成之后的总的性能(例如最后一个 epoch 的参考如下):
...Test: [40/66] Time 0.272 ( 0.285) Loss 3.0449e+00 (2.9997e+00) Acc@1 36.85 ( 36.49) Acc@5 59.11 ( 60.83)Test: [40/66] Time 0.272 ( 0.285) Loss 3.0449e+00 (2.9997e+00) Acc@1 36.85 ( 36.49) Acc@5 59.11 ( 60.83)Test: [50/66] Time 0.275 ( 0.283) Loss 3.1013e+00 (3.0017e+00) Acc@1 34.11 ( 36.38) Acc@5 59.51 ( 60.88)Test: [50/66] Time 0.275 ( 0.283) Loss 3.1013e+00 (3.0017e+00) Acc@1 34.11 ( 36.38) Acc@5 59.51 ( 60.88)Test: [50/66] Time 0.275 ( 0.283) Loss 3.1013e+00 (3.0017e+00) Acc@1 34.11 ( 36.38) Acc@5 59.51 ( 60.88)Test: [50/66] Time 0.275 ( 0.284) Loss 3.1013e+00 (3.0017e+00) Acc@1 34.11 ( 36.38) Acc@5 59.51 ( 60.88)Test: [50/66] Time 0.275 ( 0.284) Loss 3.1013e+00 (3.0017e+00) Acc@1 34.11 ( 36.38) Acc@5 59.51 ( 60.88)Test: [50/66] Time 0.275 ( 0.284) Loss 3.1013e+00 (3.0017e+00) Acc@1 34.11 ( 36.38) Acc@5 59.51 ( 60.88)Test: [50/66] Time 0.275 ( 0.284) Loss 3.1013e+00 (3.0017e+00) Acc@1 34.11 ( 36.38) Acc@5 59.51 ( 60.88)Test: [50/66] Time 0.275 ( 0.283) Loss 3.1013e+00 (3.0017e+00) Acc@1 34.11 ( 36.38) Acc@5 59.51 ( 60.88)Test: [60/66] Time 0.273 ( 0.282) Loss 3.0019e+00 (3.0007e+00) Acc@1 37.11 ( 36.44) Acc@5 61.20 ( 60.89)Test: [60/66] Time 0.273 ( 0.282) Loss 3.0019e+00 (3.0007e+00) Acc@1 37.11 ( 36.44) Acc@5 61.20 ( 60.89)Test: [60/66] Time 0.273 ( 0.282) Loss 3.0019e+00 (3.0007e+00) Acc@1 37.11 ( 36.44) Acc@5 61.20 ( 60.89)Test: [60/66] Time 0.273 ( 0.282) Loss 3.0019e+00 (3.0007e+00) Acc@1 37.11 ( 36.44) Acc@5 61.20 ( 60.89)Test: [60/66] Time 0.273 ( 0.282) Loss 3.0019e+00 (3.0007e+00) Acc@1 37.11 ( 36.44) Acc@5 61.20 ( 60.89)Test: [60/66] Time 0.273 ( 0.283) Loss 3.0019e+00 (3.0007e+00) Acc@1 37.11 ( 36.44) Acc@5 61.20 ( 60.89)Test: [60/66] Time 0.273 ( 0.282) Loss 3.0019e+00 (3.0007e+00) Acc@1 37.11 ( 36.44) Acc@5 61.20 ( 60.89)Test: [60/66] Time 0.273 ( 0.282) Loss 3.0019e+00 (3.0007e+00) Acc@1 37.11 ( 36.44) Acc@5 61.20 ( 60.89)* Acc@1 36.486 Acc@5 60.978* Acc@1 36.486 Acc@5 60.978* Acc@1 36.486 Acc@5 60.978* Acc@1 36.486 Acc@5 60.978* Acc@1 36.486 Acc@5 60.978* Acc@1 36.486 Acc@5 60.978* Acc@1 36.486 Acc@5 60.978* Acc@1 36.486 Acc@5 60.978
同时目录下会出现 model_best.pth.tar 和 checkpoint.pth.tar,记录模型权重。可以看到测试时同一个 iteration 有多条记录,这是因为 DDP 在每一张卡上都有一个进程,每一张卡都会打印一条记录。
参考文档
ViT