Create a function/script to create different models for different purposes (e.g. big, medium, small)

[mrdbourke]

I'd like a script to create 3 different models (or more):

1. Big - can be used for labelling/label improvements, this model will be the biggest and best but will be almost impossible to deploy properly, see here: https://huggingface.co/BAAI/EVA (this model requires a GPU with at least 24GB for batches of 16)
2. Medium - can be used for experimenting, reviewing labels, etc
3. Small - smallest and fastest model used for deploying (this model won't be as accurate as the others but will definitely be the fastest/capable of running on smaller devices/in the browser), see here: https://pytorch.org/vision/stable/models/generated/torchvision.models.quantization.shufflenet_v2_x2_0.html#shufflenet-v2-x2-0

Ideas

I could have a Python script called create_models.py which contains 3x functions (one for each model).

Then in the train.py script, I can call the appropriate model for the job.

Could use the "strategy" pattern in Python, see this video: https://youtu.be/WQ8bNdxREHU