스크립트

github 루트 폴더에 포함된 학습, 검증, 추론 및 체크포인트 정리 스크립트입니다. 현재 스크립트는 pip 릴리스에 패키징되어 있지 않습니다.

학습 및 검증 스크립트는 PyTorch Imagenet Examples의 초기 버전에서 발전했습니다. 시간이 지남에 따라 NVIDIA’s APEX Examples에 기반한 CUDA 전용 성능 향상 기능 등 상당한 기능을 추가했습니다.

학습 스크립트

학습 인자의 종류가 다양하며, 모든 옵션 조합(또는 옵션 자체)이 완전히 테스트된 것은 아닙니다. 학습 데이터셋 폴더의 경우, train 폴더와 validation 폴더를 포함하는 기본 폴더를 지정하세요.

단일 노드 실행의 경우, GPU당 하나의 프로세스로 torchrun을 사용하세요. 예를 들어, MobileNetV4 config를 mobilenetv4.yaml로 저장한 후:

torchrun --standalone --nproc_per_node=2 train.py \
  --config mobilenetv4.yaml --data-dir /data/imagenet --experiment mobilenetv4-cm

./distributed_train.sh 2 ...은 torchrun을 감싼 편의 래퍼입니다. GPU가 하나인 경우 python train.py ...를 사용하고 HParams에 설명된 대로 배치 크기와 학습률을 조정하세요.

`--amp` uses PyTorch native mixed precision, with float16 by default. Use `--amp-dtype bfloat16` on supported hardware. The examples below use native PyTorch AMP and DDP; Apex is not required.

설정 파일

모든 학습 인수는 --config를 통해 YAML 파일로 전달할 수 있으며, 인수는 밑줄이 포함된 이름으로 키를 지정합니다(--weight-decay는 weight_decay가 됩니다). 커맨드 라인 인수는 파일의 값을 덮어씁니다. 각 실행은 유효한 인수의 전체 집합을 출력 폴더의 args.yaml에 기록하며, 이 파일 자체가 유효한 --config이므로 python train.py --config output/train/my-run/args.yaml --experiment my-run-2로 실행을 재현할 수 있습니다. 이 파일에는 output, experiment, resume도 기록되므로, 새로운 실험 이름이 없으면 명령이 원래 폴더에 기록되고, --resume auto를 사용하면 실행을 반복하는 대신 원래 실행을 계속합니다. 설정 키는 스크립트의 인수를 기준으로 검증됩니다: 알려지지 않은 키는 경고를, 숫자 인수의 잘못된 선택지 또는 숫자가 아닌 값은 오류를 발생시킵니다.

재개 및 마일스톤 체크포인트

각 실행은 last.pth.tar(항상 최신 epoch이며 모델, 옵티마이저, EMA 및 AMP 스케일러 상태를 포함), model_best.pth.tar, 그리고 --checkpoint-hist로 제한되는 checkpoint-N.pth.tar 파일의 순위별 히스토리를 작성합니다.

• --resume PATH은 해당 상태를 복원하고 체크포인트의 epoch부터 스케줄을 계속합니다. 이는 기존 체크포인트에서 새로운 실험 폴더를 시작하는 일반적인 방법입니다.
• --resume auto는 해당 위치에서 재개합니다. --output와 --experiment가 설정된 경우, 해당 실험 폴더에 last.pth.tar가 존재하면 이를 기반으로 계속 실행하고, 그렇지 않으면 새로운 실행을 시작합니다. 중단된 후에도 동일한 명령을 다시 실행할 수 있습니다. 체크포인트 히스토리, 최적 지표 추적 및 summary.csv가 모두 일관되게 유지되므로, 재시작을 거치더라도 --checkpoint-hist의 제거(eviction)와 model_best가 정확히 유지됩니다.
• --stop-epoch N은 --epochs로 설정된 스케줄을 변경하지 않고 N개의 epoch가 완료된 후 학습을 중단합니다. 마지막 epoch는 항상 평가되고 체크포인트가 저장되므로, 학습을 세그먼트 단위로 진행할 수 있습니다.
• --save-epochs N [N ...]는 랭크된 히스토리와 별도로 불변의 milestone-N.pth.tar 스냅샷을 유지합니다(N은 완료된 에포크 수입니다). 파일 시스템이 허용하는 경우 last.pth.tar에서 하드링크로 생성되어 추가적인 저장 공간을 차지하지 않으며, 해당 에포크에서 평가가 강제됩니다.

이 두 가지가 결합되면, 안정된 트렁크를 한 번 학습한 후 그 마일스톤에서 옵티마이저 상태를 완전히 포함하여 독립적인 쿨다운을 분기하는 방식으로, 수동으로 warmup-stable-decay(WSD) 실험을 쉽게 실행할 수 있습니다.

# stable trunk through 30 completed epochs, keeping milestones at 18, 24 and 30. The 1 epoch decay is never reached.
python train.py --config configs/vit.yaml --sched wsd --sched-on-updates --epochs 31 --decay-epochs 1 \
    --stop-epoch 30 --save-epochs 18 24 30 --output output/wsd --experiment stable --resume auto

# a 6 epoch 1-sqrt cooldown branched from epoch 24, i.e. a 30 epoch run that shares the first 24 epochs with the trunk
python train.py --config configs/vit.yaml --sched wsd --sched-on-updates --epochs 30 --decay-epochs 6 --lr-decay-fn 1-sqrt \
    --resume output/wsd/stable/milestone-24.pth.tar --output output/wsd --experiment decay-p24-d6

감쇠는 스케줄의 마지막 --decay-epochs 에포크를 차지하므로, P 에포크 완료 후 D 에포크의 쿨다운은 --epochs P+D --decay-epochs D 입니다(--warmup-prefix를 사용하는 경우 --epochs에서 --warmup-epochs를 뺍니다). 아래 sweep.py는 이를 자동화합니다.

하이퍼파라미터 스윕

sweep.py(PyTorch 2.5 이상)는 train.py의 모든 인자에 대해 WSD 분기를 선택적으로 포함하여 스윕을 계획하고 실행합니다. 각 작업마다 완전한 train.py --config 파일 하나와 작업 의존성 그래프를 담은 plan.json을 작성한 후, 로컬 GPU 풀에서 실행 가능한 작업을 실행합니다. 워커가 실행되는 동안 계획은 고정됩니다. 모든 작업은 자체 실험 폴더를 가진 일반적인 train.py 호출이며, 생성된 각 run.sh는 수동으로 실행하거나 다른 스케줄러에 전달할 수 있습니다.

빠른 스윕에는 spec 파일이 필요하지 않습니다. 검색 축은 --sweep에서, 기본 구성은 -- 이후의 train.py 인자에서 가져옵니다:

python sweep.py run --name lr-wd --sweep lr=1e-4,3e-4,1e-3 --sweep weight_decay=0.05,0.1 --gpus 0,1,2,3 \
    -- --model vit_small_patch16_224 --data-dir /data/imagenet --sched cosine --epochs 100 --amp

--gpus는 실행할 GPU를 CUDA_VISIBLE_DEVICES가 설정된 경우 해당 값에 대한 인덱스로 나열합니다(예: Slurm 할당).

WSD 스윕은 여러 쿨다운에 걸쳐 하이퍼파라미터 조합별로 하나의 안정된 트렁크를 공유합니다. 아래 각 학습률에 대해, 트렁크는 18, 24, 30에 마일스톤을 두고 30에포크까지 실행되며, 각 마일스톤에서 6에포크의 쿨다운이 분기됩니다. 따라서 전체 실행 12회 대신 트렁크 4개와 쿨다운 12개가 생성됩니다. 트렁크가 계속 실행되는 동안, 각 쿨다운은 해당 마일스톤이 존재하는 즉시 시작됩니다. GPU가 하나 이상 필요한 작업은 torchrun으로 시작됩니다:

python sweep.py run --name vit-wsd --sweep lr=1e-3,5e-4,2e-4,1e-4 \
    --wsd-points 18 24 30 --wsd-decay-epochs 6 --wsd-decay-fn 1-sqrt --gpus-per-job 2 --gpus 0,1 \
    -- --model vit_small_patch16_224 --data-dir /data/imagenet --warmup-epochs 3 --warmup-prefix --amp

planner가 다음 플래그를 소유하므로, 커맨드 라인에서 지정하거나 sweep할 수 없습니다.

모든 스윕에서:

• --output
• --experiment
• --resume
• --stop-epoch
• --save-epochs

추가로, WSD 스윕에서:

• --sched (플래너가 어쨌든 설정하므로 wsd로 지정할 수 있습니다)
• --epochs
• --decay-epochs
• --lr-decay-fn

YAML 스펙은 분포를 활용한 랜덤 서치, 페어링(지핑)된 축, 시드 반복, 기본 설정 파일을 추가하며, 스키마는 sweep.py 상단의 docstring을 참고하세요. python sweep.py plan spec.yaml은 아무것도 실행하지 않고 설정 파일과 graph.dot를 작성합니다.

sweep는 spec 파일(명령줄 sweep의 경우 작업 디렉터리)을 기준으로 output/sweeps/<name>/에 저장되며, trial마다 하나의 폴더가 생성됩니다. 폴더 이름은 sweep된 값(trial-000-lr0.001)을 따르며, 각 폴더에는 job 실험 폴더가 포함됩니다. 유지되는 모든 checkpoint는 옵티마이저 상태를 포함한 전체 복사본이므로, sweep job은 last.pth.tar와 model_best.pth.tar 외에 랭킹이 매겨진 checkpoint 하나를 유지합니다(--checkpoint-hist 1). 또한 WSD sweep는 업데이트마다 스케줄을 진행합니다(--sched-on-updates). 기본 구성에서 이 중 하나를 설정하면 해당 설정이 오버라이드됩니다. 실행 중이거나 실행 후:

python sweep.py status output/sweeps/vit-wsd    # job states and progress
python sweep.py results output/sweeps/vit-wsd   # ranked endpoint metrics, written to results.csv

결과는 동일한 예산에서 실행된 트라이얼을 비교합니다. WSD 스윕은 트라이얼마다 한 행, 쿨다운마다 한 열(p24-d6은 24에포크에서 6에포크 쿨다운을 의미)을 가지며, 반복 실행의 평균으로 계산된 최종 메트릭을 표시하고 각 열의 최고값을 표시합니다. 일반 스윕은 각 에포크 예산 내에서 최종 메트릭 기준으로 순위가 매겨집니다. results.csv에는 트렁크 세그먼트를 포함하여 모든 작업과 각각의 최고 에포크가 저장됩니다. 검증 세트가 없으면 순위 매김은 학습 손실로 대체되며, 엔드포인트 체크포인트는 validate.py로 평가할 수 있습니다.

중단된 sweep에 대해 sweep.py run을 다시 실행하면 sweep가 재개됩니다. 완료된 작업은 건너뛰고, 부분적으로 실행된 작업은 해당 작업의 last.pth.tar에서 이어서 실행됩니다. sweep 폴더를 공유하는 여러 머신의 워커가 각각 실행할 수 있습니다. state/locks 아래에 있는 flock 잠금으로 조정하며, 각 작업의 모든 학습 프로세스도 이 잠금을 보유합니다. 따라서 워커가 죽은 작업은 마지막 프로세스가 종료될 때까지 다시 시작되거나 재계획되지 않습니다. 여러 머신이 공유하는 sweep에는 NFSv4 또는 -o flock 옵션으로 마운트된 Lustre처럼 모든 머신 간에 flock이 작동하는 파일 시스템이 필요합니다. 로컬에서만 잠금이 작동하는 마운트는 머신 간에 조정을 하지 않으므로, 각 클라이언트에서 잠금 상태를 확인해야 합니다. 워커를 중지하면(Ctrl-C 또는 scancel에서 오는 SIGTERM) 해당 워커의 작업이 중지되고, 작업의 run.sh를 수동으로 실행하면 조정 메커니즘을 우회합니다.

state/locks 아래에 있는 잠금 파일은 영구적입니다: 바쁜 sweep를 해제하려고 삭제하거나 교체하지 마세요. 대신 해당 워커와 학습 프로세스를 중지하거나 완료될 때까지 기다리세요. running (gone) 상태는 학습 잠금이 해제되었지만 해당 시도가 조정되지 않았음을 의미합니다. 워커를 실행하여 이를 해결하고 완료되지 않은 작업을 재개하세요. status와 results는 작업 상태를 복구하지 않습니다.

워커는 시작 시 계획을 한 번만 읽습니다. 스윕을 변경하려면 해당 스윕의 모든 워커와 학습 프로세스를 중지한 뒤, --force 플래그를 사용하여 계획을 다시 생성하고, 다시 시작해야 합니다. 어떤 작업도 시작되기 전에는 계획의 모든 내용을 변경할 수 있습니다. 시작 후에는 기존 학습 작업은 확장만 가능합니다. 기존 작업 뒤에 더 많은 트라이얼, 쿨다운 또는 WSD 포인트를 추가할 수 있으며, 각 트렁크는 중단된 위치에서 계속 진행합니다. 기존 학습 작업에 대한 변경 사항은 새로운 스윕이 필요합니다. 생성된 경로는 계획이 작성될 때 고정되므로, 스윕 폴더를 이동하는 대신 새로운 위치에 계획을 다시 생성해야 합니다.

계획 작성이 중단되면, 워커는 불완전한 계획 작성에서 설정 파일이나 런처를 감지하고 시작을 거부합니다. 원래 계획 명령을 --force 플래그와 함께 다시 실행하여 계획을 완료하거나 복구한 후, 워커를 다시 시작합니다. 각 성공적인 계획 작성은 plan.json, 설정 파일 및 런처에 해당하는 생성물을 게시합니다. 변경되지 않고 완전한 계획은 아무것도 작성하지 않습니다.

예를 들어, 위의 CLI WSD sweep를 36개의 안정 에포크로 확장하려면 워커와 작업을 중지한 다음, 원래의 계획 인자에 후속 지점을 추가하여 반복합니다:

python sweep.py plan --force --name vit-wsd --sweep lr=1e-3,5e-4,2e-4,1e-4 \
    --wsd-points 18 24 30 36 --wsd-decay-epochs 6 --wsd-decay-fn 1-sqrt --gpus-per-job 2 \
    -- --model vit_small_patch16_224 --data-dir /data/imagenet --warmup-epochs 3 --warmup-prefix --amp
python sweep.py run output/sweeps/vit-wsd --gpus 0,1

YAML sweep의 경우, 원래 spec에 지점을 추가한 후 재시작 전에 python sweep.py plan spec.yaml --force를 실행하세요. wsd.decay_epochs은 명시적으로 유지하고 변경하지 마세요(이 예시에서는 [6]): 기본값은 최종 지점의 20%이므로, 지평선을 확장하면 기존 cooldown이 변경될 수 있습니다. 기존 trial ID를 유지하세요. 그리드 축을 재배치하거나 확장하면 trial이 다시 번호가 매겨질 수 있습니다. lr_noise를 사용하는 경우, trunk의 스케줄은 지평선에 의존하므로, 이미 시작된 sweep에 후속 지점을 추가하려면 새로운 sweep가 필요합니다.

학습 설정은 플랜에 포함되어야 합니다: 스펙의 train 블록을 사용하거나, CLI 스윕을 생성할 때 -- 뒤에 오는 인수를 사용하세요. 런타임 학습 오버라이드는 없습니다. 랭킹 메트릭과 같은 보고 설정 및 출력 설정(checkpoint_hist, log_wandb, wandb_project, wandb_tags)은 --force를 사용하여 오프라인에서 변경할 수 있으며, 이 경우 작업 상태를 초기화하거나 학습 ID를 변경하지 않습니다. 출력 설정은 향후 시도에만 적용되며, 완료된 작업은 다시 실행되지 않습니다.

sweep.py는 Unix flock을 요구합니다. Linux에서 테스트되었으며, macOS에서는 작동할 것으로 예상되지만 테스트되지 않았습니다. 네이티브 Windows는 지원되지 않으며, WSL2를 사용해야 합니다. NFS 리스보다 오래 지속되는 네트워크 분할에서는 잠금이 유지되지 않습니다. 연결이 끊긴 프로세스가 잠금을 잃은 후에도 학습을 계속할 수 있으므로, 다른 머신에서 겹치는 시도가 시작될 수 있습니다.

검증 / 추론 스크립트

Validation 및 추론 스크립트는 사용법이 유사합니다. 하나는 검증 세트에 대한 메트릭을 출력하고, 다른 하나는 csv로 topk 클래스 id를 출력합니다. 학습 스크립트에서와 같이 base 폴더가 아니라, 검증 이미지가 포함된 폴더를 지정하세요.

모델의 사전 학습된 가중치(존재하는 경우)로 검증하려면:

python validate.py --data-dir /imagenet/validation/ --model seresnext26_32x4d --pretrained

checkpoint에서 추론을 실행하려면:

python inference.py --data-dir /imagenet/validation/ --model mobilenetv3_large_100 --checkpoint ./output/train/model_best.pth.tar

학습 예시

ImageNet 학습 레시피

이 명령어는 선택된 HParams 레시피에 해당하며, 여기서 복사 가능한 YAML 설정, 소스 gist, 공개된 체크포인트 결과를 찾을 수 있습니다. /data/imagenet/train/ 및 /data/imagenet/validation/ 아래에 ImageFolder 데이터가 있다고 가정합니다. --batch-size는 GPU당 값입니다. MobileNetV4 및 ConvNeXt 예시에 명시된 학습률은 주어진 전체 배치 크기에 적용되며, GPU 수가 변경되어도 자동으로 조정되지 않습니다.

MobileNetV4 Conv Medium, ImageNet-1k

mnv4_cm_r256_e500_gpu2.yaml에서: GPU 2개, 전역 배치 1024, 256px 학습, 그리고 20 warmup + 500 cosine epochs. 해당 가중치는 mobilenetv4_conv_medium.e500_r256_in1k입니다.

torchrun --standalone --nproc_per_node=2 train.py --data-dir /data/imagenet \
  --model mobilenetv4_conv_medium.e500_r256_in1k --img-size 256 --crop-pct 0.95 --batch-size 512 \
  --opt adamw --lr 0.002 --weight-decay 0.1 --clip-grad 5.0 \
  --sched cosine --sched-on-updates --epochs 500 --warmup-epochs 20 --warmup-prefix \
  --warmup-lr 0 --min-lr 0 --cooldown-epochs 0 \
  --aa rand-m8-inc1-mstd1.0 --mixup 0.2 --mixup-prob 0.8 --smoothing 0.1 \
  --reprob 0.25 --remode pixel --drop 0.2 --drop-path 0.1 \
  --model-ema --model-ema-decay 0.9998 --model-ema-warmup --dist-bn reduce \
  --amp --amp-dtype bfloat16 --torchcompile inductor --experiment mobilenetv4-cm

GPU가 bfloat16을 지원하지 않는 경우 --amp-dtype float16을 사용하세요. --torchcompile inductor은 선택 사항입니다.

레지스터가 포함된 ViT Medium, ImageNet-1k

train-vit_medium_reg4-4gpu-in1k.yaml에서: GPU 4개, 글로벌 배치 1024, 웜업 + 10쿨다운 에포크를 포함한 600 에포크. 해당 가중치는 vit_medium_patch16_reg4_gap_256.sbb_in1k입니다.

torchrun --standalone --nproc_per_node=4 train.py --data-dir /data/imagenet \
  --model vit_medium_patch16_reg4_gap_256.sbb_in1k --model-kwargs fix_init=True --head-init-scale 0 \
  --img-size 256 --crop-pct 0.95 --batch-size 256 \
  --opt adamw --opt-eps 1e-8 --lr-base 0.0008 --lr-base-size 512 --lr-base-scale sqrt \
  --weight-decay 0.07 --clip-grad 5.0 --sched cosine --sched-on-updates \
  --epochs 600 --warmup-epochs 20 --warmup-lr 5e-7 --min-lr 5e-7 --cooldown-epochs 10 \
  --aa rand-m7-inc1-mstd1.0-n3 --aug-repeats 3 --color-jitter 0 \
  --mixup 0.8 --cutmix 1.0 --smoothing 0.1 --reprob 0.2 --remode pixel --drop-path 0.2 \
  --seed 0 --amp --amp-dtype float16 --torchcompile inductor --experiment vit-medium-sbb

이 실행은 EMA 없이 AdamW를 사용합니다. 피크 LR은 글로벌 배치 1024에서 약 0.001131입니다.

ConvNeXt Tiny, ImageNet-12k에서 ImageNet-1k으로 fine-tuning

hparams_84_00-convnext_tiny.in12k-args.yaml에서: GPU 8개, 글로벌 배치 4096, 워밍업 + 10쿨다운 에포크를 포함한 75에포크. ImageNet-12k 가중치를 로드하고 --num-classes 1000을 설정하여 분류기를 교체하고 파인튜닝을 수행합니다. 결과적으로 생성되는 체크포인트 계열은 convnext_tiny.in12k_ft_in1k입니다.

torchrun --standalone --nproc_per_node=8 train.py --data-dir /data/imagenet \
  --model convnext_tiny.in12k --pretrained --num-classes 1000 \
  --img-size 224 --crop-pct 1.0 --batch-size 512 --grad-checkpointing \
  --opt adamw --lr 1e-4 --weight-decay 0.05 --layer-decay 0.75 --clip-grad 3.0 \
  --sched cosine --epochs 75 --warmup-epochs 10 --warmup-lr 1e-6 --min-lr 5e-7 \
  --cooldown-epochs 10 --lr-noise 0.1 1.0 \
  --aa rand-m8-inc1-mstd101 --mixup 0.8 --cutmix 1.0 --smoothing 0.1 \
  --reprob 0.3 --remode pixel --drop-path 0.1 --model-ema --model-ema-decay 0.9998 \
  --amp --amp-dtype float16 --experiment convnext-tiny-ft

멀티 라벨 분류

여러 개의 독립된 라벨을 가진 이미지의 경우 --task multilabel을 사용하세요.

• --num-classes는 어휘 크기를 설정합니다.
• --target-format indices(기본값)은 대상 컬럼을 양의 0 기반 클래스 ID 목록으로 읽습니다. 레이블은 배치 처리 전에 float multi-hot target으로 변환되며, 빈 목록은 양의 레이블이 없는 이미지를 나타냅니다.
• --target-format multihot는 --num-classes 길이의 dense multi-hot 벡터([0, 1] 범위의 bool, int 또는 float 값)를 읽습니다. 이 벡터는 그대로 사용되므로, soft 값은 BCE로 직접 학습되고 평가 지표 계산 시 0.5 기준으로 이진화됩니다. multi-hot 대상의 class map은 source class index를 키로 사용해야 합니다.

flag가 없으면 정수형 멀티핫 벡터는 인덱스 리스트와 구별할 수 없으므로 모든 리스트가 클래스 인덱스로 읽힙니다.

plant pathology dataset은 다음과 같은 형식을 제공합니다: image 열과 ClassLabel 값의 시퀀스를 포함하는 labels 열입니다. 여섯 개의 클래스는 complex, frog_eye_leaf_spot, healthy, powdery_mildew, rust, scab입니다. 예를 들어, [5, 1, 0]은 [1, 1, 0, 0, 0, 1]로 변환됩니다.

python train.py --dataset hfds/timm/plant-pathology-2021 \
  --task multilabel --target-key labels --num-classes 6 \
  --model resnet50 --pretrained --batch-size 32 --epochs 30 \
  --opt adamw --lr 1e-4 --weight-decay 0.01 --amp

python validate.py --dataset hfds/timm/plant-pathology-2021 \
  --task multilabel --target-key labels --num-classes 6 \
  --model resnet50 --checkpoint ./output/train/<experiment>/model_best.pth.tar

hfids/timm/plant-pathology-2021은 동일한 대상 형식으로 스트리밍을 지원합니다. 리더 지원:

• Hugging Face hfds와 hfids는 다중 레이블 타깃을 지원합니다.
• WebDataset (wds)는 각 샘플의 .json 사이드카 파일에 --target-key로 지정된 필드 아래에 저장된 클래스 ID 목록을 지원합니다. --target-key가 없으면 .cls 파일이 단일 클래스 인덱스를 제공합니다.
• Folder and tar readers는 단일 클래스 인덱스를 반환합니다.

각 라벨마다 하나의 바이너리 필드를 저장하는 Wake Vision과 같은 데이터셋은 다음 설정을 사용합니다:

• --target-key person,depiction,body_part은 쉼표로 구분된 필드 목록을 선택합니다. 각 필드는 나열된 순서대로 하나의 클래스가 되며, 0보다 큰 값은 양수로 계산됩니다. 이 기능은 hfds, hfids, tfds 리더에서 동작하며, 필드가 .json 사이드카에 있는 경우 WebDataset에서도 동작합니다.
• --num-classes은 필드 수와 같아야 합니다.
• 멀티 필드 키는 자체 포맷을 정의하므로 --target-format은 기본값을 유지합니다.
• 중첩 필드 경로에는 /를 사용합니다. 예를 들어 TFDS bee_dataset의 경우 --input-key input --target-key output/pollen_output,output/wasps_output와 같이 사용할 수 있습니다. 단일 대상 키도 동일한 경로를 허용합니다. /를 포함하는 리터럴 필드 이름(예: TFDS image/filename)은 경로보다 우선합니다.

Standard 및 NaFlex 로더는 mixup과 cutmix를 포함한 다중 레이블 타깃을 지원합니다. NaFlex 학습은 현재 hfds와 같은 map 스타일 데이터셋이 필요합니다.

이 작업은 logits를 사용하는 binary cross-entropy로 학습합니다. 다중 레이블 학습에서 Smoothing 기본값은 0입니다. --smoothing s를 설정하면 각 binary target이 target * (1 - s) + 0.5 * s로 이동하며, Mixup/CutMix가 활성화된 경우 해당 기법에 의해, 그렇지 않은 경우 손실 함수에 의해 적용됩니다. 기존 --bce-sum 및 --bce-target-thresh 옵션도 적용됩니다. 이 작업은 Distillation 및 JSD loss를 지원하지 않습니다.

--loss는 다른 손실 함수를 선택합니다(bce는 멀티 라벨 기본값입니다):

• asl: Asymmetric Loss. BCE와 동일한 smoothing 및 Mixup/CutMix 동작을 따르며, float32로 계산되므로 --amp와 함께 안전하게 사용할 수 있습니다.

  • --asl-gamma-neg는 음의 대상에 대한 포커싱을 설정합니다(기본값 4).
  • --asl-gamma-pos는 양의 대상에 대한 포커싱을 설정합니다(기본값 1).
  • --asl-clip는 음수(negative)의 확률 마진을 설정합니다(기본값 0.05).
  • --asl-reduction의 기본값은 batchmean이며, 클래스에 대해 손실을 합산하고 배치에 대해 평균을 냅니다. sum은 배치에 대해서도 합산하여 유효 학습률이 배치 크기에 따라 스케일링되도록 하는 원래의 ASL reduction을 복원합니다.

• twoway: Two-way Multi-Label Loss, 각 샘플의 클래스와 각 클래스의 샘플 간에 양의 로짓과 음의 로짓을 대비합니다. 로짓은 랭킹 점수이므로 mAP가 의미 있는 지표이며, 시그모이드 임계값에서의 F1은 보정되지 않습니다.

  • --twoway-tp는 양수 logit의 temperature를 설정합니다(기본값 4).
  • --twoway-tn은 음의 로짓 온도(기본값 1)를 설정합니다.

• zlpr: ZLPR, 0을 내장된 logit 임계값으로 사용하는 softmax cross-entropy의 하이퍼파라미터가 없는 다중 레이블 확장입니다.

• db: 긴 꼬리(long-tailed) 다중 레이블 데이터용 Distribution-Balanced Loss. 레이블 공출현(co-occurrence)을 고려한 샘플별 재균형화 가중치를 적용한 BCE, 로짓에 대한 음수 허용 정규화, 포칼 항을 포함합니다. Mixup/CutMix 및 스무딩을 지원하며, 기본값은 참조 DB 구성에서 가져옵니다:

  • --db-neg-scale는 음수 로짓 스케일(기본값 2)을 설정합니다.
  • --db-init-bias은 클래스 선행 로짓 시프트 계수를 설정합니다(기본값 0.05).
  • --db-focal-gamma는 focal focusing을 설정합니다(기본값 2).
  • --class-stats는 필요한 학습 레이블 개수를 클래스 통계 파일(아래)로 제공합니다.

twoway와 zlpr는 바이너리 타겟이 필요하므로 Mixup/CutMix 또는 --smoothing를 지원하지 않습니다.

불균형한 라벨 분포의 경우, bce와 asl은 두 가지 가중치 플래그를 지원합니다:

• --loss-pos-weight는 각 클래스의 양(positive) 항을 가중하여 클래스 내에서 양과 음(negative)을 균형 있게 조정합니다. 이는 예측을 양의 방향으로 이동시키며(고정된 임계값에서 recall이 높아지고, 확률의 보정이 덜 이루어짐)
• --loss-class-weight는 각 클래스의 전체 손실을 가중하여 클래스 간 균형을 맞추고 보정을 유지합니다.

각 플래그은 명시적인 가중치(단일 값, 클래스당 하나의 값을 가진 쉼표로 구분된 목록, 또는 파일) 또는 --class-stats로 제공된 학습 레이블 개수에서 가중치를 유도하는 메서드 이름을 받습니다:

• --loss-pos-weight neg_pos: 클래스당 negatives / positives 비율.
• --loss-class-weight inv_freq: 역 클래스 빈도, 평균이 1이 되도록 정규화됩니다.
• --loss-class-weight effective_num: 샘플의 유효 개수에서 계산한 class-balanced 가중치로, 평균이 1이 되도록 정규화됩니다.

파생 가중치 옵션:

• --loss-weight-power은 neg_pos와 inv_freq의 지수를 설정합니다(기본값 1). 0.5를 사용하면 더 부드러운 제곱근 재가중화를 적용할 수 있습니다.
• --loss-weight-beta는 effective_num 가중치에 대한 beta를 설정합니다(기본값 0.999).
• --loss-weight-max은 유도된 가중치를 상한으로 제한합니다.

두 가중치 플래그 모두 동일한 통계 정보를 사용할 수 있습니다. class_weights.py는 train.py와 동일한 데이터셋 인수를 사용하여 분할의 라벨을 집계하고, 클래스 통계 파일(--method은 해당 방법의 가중치를 미리 보여 줌)을 작성합니다:

python class_weights.py --dataset hfds/timm/plant-pathology-2021 --split train \
  --task multilabel --target-key labels --num-classes 6 --output stats.json

python train.py --dataset hfds/timm/plant-pathology-2021 \
  --task multilabel --target-key labels --num-classes 6 \
  --class-stats stats.json --loss-pos-weight neg_pos --loss-weight-power 0.5 ...

비추천 플래그:

Validation은 BCE loss, map(mean average precision), micro_f1, macro_f1, sample_f1를 보고합니다. 지표는 백분율이며, 분산 랭크에 걸쳐 수집된 예측을 포함한 전체 validation set을 사용합니다. 정렬된 분산 sampler padding은 제외됩니다. Exact mAP는 평가가 완료될 때까지 예측을 CPU에 유지하므로, 메모리 사용량은 예시와 클래스의 수에 따라 증가합니다. 양성이 없는 클래스는 AP에 0을 기여하며, 분모가 0인 F1 값은 0입니다.

--eval-metric map는 체크포인트 선택 및 스케줄링의 기본값입니다. F1은 0.5의 시그모이드 확률 임계값을 사용하며, --multilabel-threshold로 설정할 수 있습니다. mAP는 이 임계값과 무관합니다. 모델 출력은 로짓(logits)으로 유지되므로, 학습된 모델을 사용할 때는 독립적인 클래스 확률을 얻기 위해 시그모이드를 적용해야 합니다.

과거의 ImageNet 레시피

이전 실행 결과는 RMSProp, SGD, AugMix/JSD 학습에 대한 유용한 참고 자료로 남아 있습니다. 하드웨어 및 보고된 지표는 원래 실행을 설명합니다. 새로운 AdamW 레시피의 경우, 위의 예제 또는 gist 컬렉션부터 시작하세요.

RandAugment을 적용한 EfficientNet-B2 - top-1 80.4, top-5 95.1

원래 실행에서는 NVIDIA Apex를 사용한 듀얼 Titan RTX 카드를 사용했으며, 현재 --amp는 네이티브 PyTorch AMP를 사용합니다:

./distributed_train.sh 2 --data-dir /imagenet/ --model efficientnet_b2 -b 128 --sched step --epochs 450 --decay-epochs 2.4 --decay-rate .97 --opt rmsproptf --opt-eps .001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.3 --drop-path 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --amp --lr .016

RandAugment을 적용한 MixNet-XL - 80.5 top-1, 94.9 top-5

원래 실행에서는 NVIDIA Apex를 사용한 듀얼 Titan RTX 카드를 사용했으며, 현재 --amp는 네이티브 PyTorch AMP를 사용합니다:

./distributed_train.sh 2 --data-dir /imagenet/ --model mixnet_xl -b 128 --sched step --epochs 450 --decay-epochs 2.4 --decay-rate .969 --opt rmsproptf --opt-eps .001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.3 --drop-path 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.3 --amp --lr .016 --dist-bn reduce

SE-ResNeXt-26-D 및 SE-ResNeXt-26-T

이러한 hparams(또는 유사한 값)은 다양한 ResNet 아키텍처에서 잘 작동하며, 일반적으로 모델 크기가 커질수록 epoch 수를 늘리는 것이 좋습니다. 예를 들어, ResNe(X)t50의 경우 약 180~200, 더 큰 모델의 경우 220 이상을 권장합니다. 더 나은 GPU를 사용하거나 AMP가 활성화된 경우, batch size와 LR을 비례하여 늘리세요. 이 파라미터는 1080Ti 카드 2개를 기준으로 설정되었습니다:

./distributed_train.sh 2 --data-dir /imagenet/ --model seresnext26t_32x4d --lr 0.1 --warmup-epochs 5 --epochs 160 --weight-decay 1e-4 --sched cosine --reprob 0.4 --remode pixel -b 112

RandAugment을 적용한 EfficientNet-B3 - 81.5 top-1, 95.7 top-5

이 모델의 학습은 위에서 언급한 EfficientNet-B2 w/ RA와 동일한 커맨드 라인으로 시작되었습니다. 약 3주간의 학습 후 프로세스가 충돌했습니다. 결과가 인상적이지 않아 몇 가지 파라미터를 조정(RE 확률 증가, rand-aug 감소, ema-decay 증가)하여 학습을 여러 번 재개했습니다. 여전히 만족스러운 결과는 나오지 않았습니다. 결국 모든 재시작에서 가장 좋은 체크포인트를 평균화했습니다. 기본 해상도/크롭에서는 결과가 평범하지만, 1.0의 전체 이미지 테스트 크롭을 사용할 때 의외로 훨씬 더 좋은 성능을 보입니다.

RandAugment을 적용한 EfficientNet-B0 - 77.7 top-1, 95.3 top-5

Michael Klachko는 더 큰 배치 크기에 맞게 조정된 B2의 커맨드 라인과 권장되는 B0 드롭아웃 비율 0.2를 사용하여 이러한 결과를 달성했습니다.

./distributed_train.sh 2 --data-dir /imagenet/ --model efficientnet_b0 -b 384 --sched step --epochs 450 --decay-epochs 2.4 --decay-rate .97 --opt rmsproptf --opt-eps .001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.2 --drop-path 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --amp --lr .048

JSD 손실과 RandAugment(클린 + 2x RA 증강)을 적용한 ResNet50 - 79.04 top-1, 94.39 top-5

1080Ti 2장으로 학습했으며, 시간이 꽤 걸렸습니다. 첫 번째 좋은 AugMix 학습 결과인 78.99보다 ImageNet 검증 결과가 약간, 통계적으로 유의미하지는 않을 정도로 더 좋습니다. 그러나 이 가중치는 ImageNetV2, ImageNet-Sketch 등의 테스트에서 더 강건합니다. 첫 번째 AugMix 실행과 달리, SplitBatchNorm을 활성화하고, clean split에서 random erasing을 비활성화했으며, 2개의 증강 경로에서 random erasing 확률을 높였습니다.

./distributed_train.sh 2 --data-dir /imagenet -b 64 --model resnet50 --sched cosine --epochs 200 --lr 0.05 --amp --remode pixel --reprob 0.6 --aug-splits 3 --aa rand-m9-mstd0.5-inc1 --resplit --split-bn --loss jsd --dist-bn reduce

RandAugment을 적용한 EfficientNet-ES (EdgeTPU-Small) - 78.066 top-1, 93.926 top-5

Andrew Lavin이 8개의 V100 카드로 학습했습니다. 모델 EMA는 사용하지 않았으며, 최종 체크포인트는 학습 중 가장 좋았던 8개 체크포인트의 평균입니다.

./distributed_train.sh 8 --data-dir /imagenet --model efficientnet_es -b 128 --sched step --epochs 450 --decay-epochs 2.4 --decay-rate .97 --opt rmsproptf --opt-eps .001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.2 --drop-path 0.2  --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --amp --lr .064

MobileNetV3-Large-100 - 75.766 top-1, 92.542 top-5

./distributed_train.sh 2 --data-dir /imagenet/ --model mobilenetv3_large_100 -b 512 --sched step --epochs 600 --decay-epochs 2.4 --decay-rate .973 --opt rmsproptf --opt-eps .001 -j 7 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.2 --drop-path 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --amp --lr .064 --lr-noise 0.42 0.9

ResNeXt-50 32x4d w/ RandAugment - 79.762 top-1, 94.60 top-5

이 파라미터는 SE-ResNeXt-50과 SK-ResNeXt-50, 그리고 아마도 101에서도 잘 작동합니다. 저는 2개의 GPU를 사용하여 유효 배치 크기당 약간 더 높은 LR(lr=0.18, GPU당 b=192)으로 학습시킨 SK-ResNeXt-50 32x4d에 이 파라미터를 사용했습니다. 아래 명령줄은 8개 GPU 학습을 위해 조정되었습니다.

./distributed_train.sh 8 --data-dir /imagenet --model resnext50_32x4d --lr 0.6 --warmup-epochs 5 --epochs 240 --weight-decay 1e-4 --sched cosine --reprob 0.4 --recount 3 --remode pixel --aa rand-m7-mstd0.5-inc1 -b 192 -j 6 --amp --dist-bn reduce