用于 PyTorch 的 EfficientNet,搭配 DALI 和 AutoAugment#
此示例展示了 DALI 自动增强的实现方式 - 最值得注意的是 AutoAugment 和 TrivialAugment - 如何在训练中使用。它展示了 EfficientNet 的训练,EfficientNet 是一种图像分类模型,最初在 EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks 中描述。
该代码基于 NVIDIA 深度学习示例 - 它已通过 DALI pipeline 扩展,支持自动增强,可以在此处找到。
与深度学习示例配置的差异#
参数的默认值已调整为 EfficientNet 训练中使用的值。
--data-backend
参数已更改为接受dali
、pytorch
或synthetic
。默认设置为dali
。--dali-device
已添加以控制某些 DALI 算子的放置位置。--augmentation
已替换为--automatic-augmentation
,现在支持disabled
、autoaugment
和trivialaugment
值。--workers
默认值已减半以适应 DALI。当使用pytorch
数据加载器时,该值会自动加倍。因此,默认值在使用两种加载器时都表现良好。该模型仅限于 EfficientNet-B0 架构。
数据后端#
此模型使用以下数据增强
对于训练
随机调整大小裁剪到目标图像大小(在本例中为 224)
从 8% 缩放到 100%
纵横比从 3/4 到 4/3
随机水平翻转
[可选:AutoAugment 或 TrivialAugment]
归一化
对于推理
缩放到目标图像大小 + 额外的尺寸边距(在本例中为 224 + 32 = 266)
中心裁剪到目标图像大小(在本例中为 224)
归一化
设置#
EfficientNet 脚本在 ImageNet 1k 上运行,ImageNet 1k 是来自 ILSVRC 挑战赛的广泛流行的图像分类数据集。
提取训练数据
mkdir train && mv ILSVRC2012_img_train.tar train/ && cd train
tar -xvf ILSVRC2012_img_train.tar && rm -f ILSVRC2012_img_train.tar
find . -name "*.tar" | while read NAME ; do mkdir -p "${NAME%.tar}"
tar -xvf "${NAME}" -C "${NAME%.tar}"; rm -f "${NAME}"; done
cd ..
提取验证数据并将图像移动到子文件夹
mkdir val && mv ILSVRC2012_img_val.tar val/ && cd val && tar -xvf ILSVRC2012_img_val.tar
wget -qO- https://raw.githubusercontent.com/soumith/imagenetloader.torch/master/valprep.sh | bash
放置 train/
和 val/
目录的目录在本文档中称为 $PATH_TO_IMAGENET
。
确保您正在使用 NVIDIA PyTorch NGC 容器,或者您已安装 DALI 和 PyTorch。
安装 NVIDIA DLLogger 和 pynvml。
运行模型#
训练#
要在单个 GPU 上运行训练,请使用 main.py
入口点
对于 FP32:
python ./main.py --batch-size 64 $PATH_TO_IMAGENET
对于 AMP:
python ./main.py --batch-size 64 --amp --static-loss-scale 128 $PATH_TO_IMAGENET
您可能需要为您的机器调整 --batch-size
参数。
您可以通过添加以下内容来更改使用的数据加载器和自动增强方案
--data-backend
:dali
|dali_proxy
|pytorch
|synthetic
,--automatic-augmentation
:disabled
|autoaugment
|trivialaugment
(最后一个仅适用于 DALI),--dali-device
:cpu
|gpu
(仅适用于 DALI)。
默认情况下,使用带有 AutoAugment 的 DALI GPU 变体(dali
和 dali_proxy
后端)。
数据后端#
dali:利用 DALI pipeline 以及 DALI 的 PyTorch 迭代器进行数据加载、预处理和增强。
dali_proxy:使用 DALI pipeline 进行预处理和增强,同时依赖 PyTorch 的数据加载器。DALI Proxy 促进数据传输到 DALI 进行处理。请参阅 PyTorch DALI 代理。
pytorch:采用原生 PyTorch 数据加载器进行数据预处理和增强。
synthetic:动态创建合成数据,这对于测试和基准测试目的非常有用。此后端消除了对实际数据集的需求,提供了一种模拟数据加载的便捷方法。
例如,要使用 TrivialAugment 在批量大小为 128 的 AMP 上运行 EfficientNet,您需要调用
python ./main.py --amp --static-loss-scale 128 --batch-size 128 --data-backend dali --automatic-augmentation trivialaugment $PATH_TO_IMAGENET
要在多个 GPU 上运行,请使用 multiproc.py
启动 main.py
入口点脚本,并将 GPU 数量作为 --nproc_per_node
参数传递。例如,要在 8 个 GPU 上使用 AMP 和带有 AutoAugment 的 DALI 运行模型,您需要调用
python ./multiproc.py --nproc_per_node 8 ./main.py --amp --static-loss-scale 128 --batch-size 128 --data-backend dali --automatic-augmentation autoaugment $PATH_TO_IMAGENET
要查看可用选项及其描述的完整列表,请使用 -h
或 --help
命令行选项,例如
python main.py -h
使用标准配置进行训练#
要以标准配置(DGX A100/DGX-1V、AMP、400 Epochs、带有 AutoAugment 的 DALI)运行训练,请调用以下命令
对于 DGX1V-16G:
python multiproc.py --nproc_per_node 8 ./main.py --amp --static-loss-scale 128 --batch-size 128 $PATH_TO_IMAGENET
对于 DGX-A100:
python multiproc.py --nproc_per_node 8 ./main.py --amp --static-loss-scale 128 --batch-size 256 $PATH_TO_IMAGENET`
基准测试#
要使用不同的数据加载器和自动增强运行训练基准测试,您可以使用以下命令,假设它们在 DGX1V-16G 上以 8 个 GPU、128 批量大小和 AMP 运行
# Adjust the following variable to control where to store the results of the benchmark runs
export RESULT_WORKSPACE=./
# synthetic benchmark
python multiproc.py --nproc_per_node 8 ./main.py --amp --static-loss-scale 128
--batch-size 128 --epochs 1 --prof 1000 --no-checkpoints
--training-only --data-backend synthetic
--workspace $RESULT_WORKSPACE
--report-file bench_report_synthetic.json $PATH_TO_IMAGENET
# DALI without automatic augmentations
python multiproc.py --nproc_per_node 8 ./main.py --amp --static-loss-scale 128
--batch-size 128 --epochs 4 --no-checkpoints --training-only
--data-backend dali --automatic-augmentation disabled
--workspace $RESULT_WORKSPACE
--report-file bench_report_dali.json $PATH_TO_IMAGENET
# DALI with AutoAugment
python multiproc.py --nproc_per_node 8 ./main.py --amp --static-loss-scale 128
--batch-size 128 --epochs 4 --no-checkpoints --training-only
--data-backend dali --automatic-augmentation autoaugment
--workspace $RESULT_WORKSPACE
--report-file bench_report_dali_aa.json $PATH_TO_IMAGENET
# DALI with TrivialAugment
python multiproc.py --nproc_per_node 8 ./main.py --amp --static-loss-scale 128
--batch-size 128 --epochs 4 --no-checkpoints --training-only
--data-backend dali --automatic-augmentation trivialaugment
--workspace $RESULT_WORKSPACE
--report-file bench_report_dali_ta.json $PATH_TO_IMAGENET
# DALI proxy with AutoAugment
python multiproc.py --nproc_per_node 8 ./main.py --amp --static-loss-scale 128
--batch-size 128 --epochs 4 --no-checkpoints --training-only
--data-backend dali_proxy --automatic-augmentation autoaugment
--workspace $RESULT_WORKSPACE
--report-file bench_report_dali_proxy_aa.json $PATH_TO_IMAGENET
# DALI proxy with TrivialAugment
python multiproc.py --nproc_per_node 8 ./main.py --amp --static-loss-scale 128
--batch-size 128 --epochs 4 --no-checkpoints --training-only
--data-backend dali_proxy --automatic-augmentation trivialaugment
--workspace $RESULT_WORKSPACE
--report-file bench_report_dali_proxy_ta.json $PATH_TO_IMAGENET
# PyTorch without automatic augmentations
python multiproc.py --nproc_per_node 8 ./main.py --amp --static-loss-scale 128
--batch-size 128 --epochs 4 --no-checkpoints --training-only
--data-backend pytorch --automatic-augmentation disabled
--workspace $RESULT_WORKSPACE
--report-file bench_report_pytorch.json $PATH_TO_IMAGENET
# PyTorch with AutoAugment:
python multiproc.py --nproc_per_node 8 ./main.py --amp --static-loss-scale 128
--batch-size 128 --epochs 4 --no-checkpoints --training-only
--data-backend pytorch --automatic-augmentation autoaugment
--workspace $RESULT_WORKSPACE
--report-file bench_report_pytorch_aa.json $PATH_TO_IMAGENET
推理#
验证在每个 epoch 完成,也可以在检查点模型上单独运行。
python ./main.py --evaluate --epochs 1 --resume <path to checkpoint>
-b <batch size> $PATH_TO_IMAGENET
要在 JPEG 图像上运行推理,您必须首先从检查点提取模型权重
python checkpoint2model.py --checkpoint-path <path to checkpoint>
--weight-path <path where weights will be stored>
然后,运行分类脚本
python classify.py --pretrained-from-file <path to weights from previous step>
--precision AMP|FP32 --image <path to JPEG image>