I think there is a minor error

[skhnha]

Hi, thanks for making your code public. It is really great work!
I ran your code, and i think there is a minor error in your code.

On line 217 in train.py
lr_scheduler = WarmupCosineAnnealingLR(optimizer=optimizer, T_cosine_max=args.epochs * IMAGENET_TRAINSET_SIZE // args.batch_size // ngpus_per_node, warmup=args.epochs/24)

I think it makes warm up the learning rate during only 5 steps, not 5 epochs. To warm up the learning rate during 5 epochs, 'args.epochs/24' should be 'args.epochs*len(train_loader)/24'.

Therefore, I modify line 217 as follows,
lr_scheduler = WarmupCosineAnnealingLR(optimizer=optimizer, T_cosine_max=args.epochs * len(train_loader), warmup=args.epochs * len(train_loader) / 24)

Thank you!

[JUGGHM]

I think you are right but I remember having fixed it long ago... Anyway