Icevision: Add how to locally save and load trained weights in the How-To

[ai-fast-track]

📓 Documentation Update

What part of documentation was unclear or wrong?
In the How-To Guide, we show how to load trained weights from torch hub.

We need also to show how to locally save and load trained weight saved on a local machine.

[elmourr]

How to save and load a trained model from the local machine

This example shows how to locally save and load trained weights using an EffecientDet model.

Check out the Quick Start Notebook to get familiar with all the steps from the training a dataset to to the point of saving the trained weights.

# Model
model = efficientdet.model(model_name="tf_efficientdet_lite0", num_classes=len(class_map), img_size=size)

# Train the model using either Fastai Learner of Pytorch-Lightning Traine

# Save trained weights in a local folder
torch.save(model.state_dict(), path_to_local_model_folder/'model.pth')


# Load locally saved weights
model_path = Path(path_to_local_model_folder)

# Maps IDs to class names.
class_map = datasets.pets.class_map()  # here an example of the PET dataset

#define the model to be loaded
model = efficientdet.model(model_name="tf_efficientdet_lite0", num_classes=len(class_map), img_size=size)
state_dict = torch.load(model_path/'model.pth', map_location=torch.device('cpu')) 
## nb: "map_location"  will put the model on cpu, optionally move to gpu if necessary by replacing 'cpu' by 'cuda'

#load the model
model.load_state_dict(state_dict)

[lgvaz]

@ai-fast-track I guess we can close this?