Distillation support for torchvision script

[rahul-tuli]

The goal of this PR is to add distillation support to our pytorch/torchvision integration
Test recipe:

distillation.yaml:

num_epochs: &num_epochs 10
lr: &lr 0.008

training_modifiers:
  - !EpochRangeModifier
    start_epoch: 0.0
    end_epoch: *num_epochs

  - !SetLearningRateModifier
    start_epoch: 0.0
    learning_rate: *lr

  - !DistillationModifier
     start_epoch: 0.0
     hardness: 0.5
     temperature: 2.0
     distill_output_keys: [0]

Test commands, (run manually):

• self distillation

sparseml.image_classification.train \ 
    --recipe distillation.yaml --pretrained True --pretrained-dataset imagenette \
    --arch-key resnet50 --dataset-path /home/rahul/datasets/imagenette/imagenette-160 \
    --batch-size 128 --opt SGD --output-dir ./training-runs/image_classification-pretrained \
    --distill-teacher self

• This command should fail (distillation recipe given but no --distill-teacher specified)

sparseml.image_classification.train \ 
    --recipe distillation.yaml --pretrained True --pretrained-dataset imagenette \
    --arch-key resnet50 --dataset-path /home/rahul/datasets/imagenette/imagenette-160 \
    --batch-size 128 --opt SGD --output-dir ./training-runs/image_classification-pretrained 

• Disable distillation

sparseml.image_classification.train \ 
    --recipe distillation.yaml --pretrained True --pretrained-dataset imagenette \
    --arch-key resnet50 --dataset-path /home/rahul/datasets/imagenette/imagenette-160 \
    --batch-size 128 --opt SGD --output-dir ./training-runs/image_classification-pretrained \
    --distill-teacher disable

• Distill a mobilenet using a resnet50 teacher from sparsezoo

sparseml.image_classification.train \ 
    --recipe distillation.yaml --pretrained True --pretrained-dataset imagenette \
    --arch-key mobilenet --dataset-path /home/rahul/datasets/imagenette/imagenette-160 \
    --batch-size 128 --opt SGD --output-dir ./training-runs/image_classification-pretrained \
    --distill-teacher zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/base-none \
    --pretrained-teacher-dataset imagenet --teacher-arch-key resnet50

[corey-nm]

changes look good, but need to call manager.update_loss (or whatever the call is) to actually use distillation loss

[rahul-tuli]

changes look good, but need to call manager.update_loss (or whatever the call is) to actually use distillation loss

Great catch, updated!

[KSGulin]

Left one comment, otherwise LGTM!

[bfineran]

Looks great @rahul-tuli

[corey-nm]

Looks like loss_update returns the new loss to use! So close 😀

[corey-nm]

🚀 LETS GOOOOOO