|
| 1 | +📂 Saving and Loading |
| 2 | +===================== |
| 3 | + |
| 4 | +In this section, we will discuss how to save a trained model, push it to the Hugging Face Hub, and load it back for later use. |
| 5 | + |
| 6 | +Saving and Sharing a Model |
| 7 | +-------------------------- |
| 8 | + |
| 9 | +Once you have trained your model, you can save it using the `.save_pretrained` method. This method saves the model configuration and weights to a directory of your choice. |
| 10 | +And, optionally, you can push the model to the Hugging Face Hub by setting the `push_to_hub` parameter to `True`. |
| 11 | + |
| 12 | +For example: |
| 13 | + |
| 14 | +.. code:: python |
| 15 | +
|
| 16 | + import segmentation_models_pytorch as smp |
| 17 | +
|
| 18 | + model = smp.Unet('resnet34', encoder_weights='imagenet') |
| 19 | +
|
| 20 | + # After training your model, save it to a directory |
| 21 | + model.save_pretrained('./my_model') |
| 22 | +
|
| 23 | + # Or saved and pushed to the Hub simultaneously |
| 24 | + model.save_pretrained('username/my-model', push_to_hub=True) |
| 25 | +
|
| 26 | +Loading Trained Model |
| 27 | +--------------------- |
| 28 | + |
| 29 | +Once your model is saved and pushed to the Hub, you can load it back using the `smp.from_pretrained` method. This method allows you to load the model weights and configuration from a directory or directly from the Hub. |
| 30 | + |
| 31 | +For example: |
| 32 | + |
| 33 | +.. code:: python |
| 34 | +
|
| 35 | + import segmentation_models_pytorch as smp |
| 36 | +
|
| 37 | + # Load the model from the local directory |
| 38 | + model = smp.from_pretrained('./my_model') |
| 39 | +
|
| 40 | + # Alternatively, load the model directly from the Hugging Face Hub |
| 41 | + model = smp.from_pretrained('username/my-model') |
| 42 | +
|
| 43 | +Saving model Metrics and Dataset Name |
| 44 | +------------------------------------- |
| 45 | + |
| 46 | +You can simply pass the `metrics` and `dataset` parameters to the `save_pretrained` method to save the model metrics and dataset name in Model Card along with the model configuration and weights. |
| 47 | + |
| 48 | +For example: |
| 49 | + |
| 50 | +.. code:: python |
| 51 | +
|
| 52 | + import segmentation_models_pytorch as smp |
| 53 | +
|
| 54 | + model = smp.Unet('resnet34', encoder_weights='imagenet') |
| 55 | +
|
| 56 | + # After training your model, save it to a directory |
| 57 | + model.save_pretrained('./my_model', metrics={'accuracy': 0.95}, dataset='my_dataset') |
| 58 | +
|
| 59 | + # Or saved and pushed to the Hub simultaneously |
| 60 | + model.save_pretrained('username/my-model', push_to_hub=True, metrics={'accuracy': 0.95}, dataset='my_dataset') |
| 61 | +
|
| 62 | +
|
| 63 | +Conclusion |
| 64 | +---------- |
| 65 | + |
| 66 | +By following these steps, you can easily save, share, and load your models, facilitating collaboration and reproducibility in your projects. Don't forget to replace the placeholders with your actual model paths and names. |
| 67 | + |
| 68 | +|colab-badge| |
| 69 | + |
| 70 | +.. |colab-badge| image:: https://colab.research.google.com/assets/colab-badge.svg |
| 71 | + :target: https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/master/examples/binary_segmentation_intro.ipynb |
| 72 | + :alt: Open In Colab |
| 73 | + |
| 74 | + |
0 commit comments