Update sparseml to support torch 2.0

[dsikka]

• Updated max version of torch to be 2.0, updated torchvision max to be 0.15.1
• Updated max torchaudio to be 2.0.1
• Tested and passed all pytorch tests
• Tested updated package versions in training using sparseml.image_classification.train. Used the ImageNette dataset and ResNet-50 model from SparseZoo

# Epoch and Learning-Rate variables
num_epochs: 3.0
init_lr: 0.0005

# quantization variables
quantization_epochs: 1.0

training_modifiers:
  - !EpochRangeModifier
    start_epoch: 0.0
    end_epoch: eval(num_epochs)

  - !LearningRateFunctionModifier
    final_lr: 0.0
    init_lr: eval(init_lr)
    lr_func: cosine
    start_epoch: 0.0
    end_epoch: eval(num_epochs)

# Phase 1 Sparse Transfer Learning / Recovery
sparse_transfer_learning_modifiers:
  - !ConstantPruningModifier
    start_epoch: 0.0
    params: __ALL_PRUNABLE__

# Phase 2 Apply quantization
sparse_quantized_transfer_learning_modifiers:
  - !QuantizationModifier
    start_epoch: eval(num_epochs - quantization_epochs)

• Trained using the sample recipe above and then exported in onnx using sparseml.image_classification.export_onnx
• Confirmed size of the onnx export model to match another model trained/exported using torch 1.13.1 (23 mb) and validated graph using Netron thanks to @KSGulin

[bfineran]

thanks @dsikka - one last thing to look for is exporting quantized models with convert_qat=True and verifying that the weights are properly compressed (ie load a quantized resnet pytorch model from the sparsezoo, export it to onnx and verify that the file size is the expected quantized one)

[dbogunowicz]

Oh, wow, this seems to be much less effort than expected no?

[dsikka]

Size of the exported model is 23 mb using torch 2.0, identical to an onnx exported resnet model trained with torch 1.13.1

thanks @dsikka - one last thing to look for is exporting quantized models with convert_qat=True and verifying that the weights are properly compressed (ie load a quantized resnet pytorch model from the sparsezoo, export it to onnx and verify that the file size is the expected quantized one)

[rahul-tuli]

LGTM after outdated comment removal!