Skip to content

Commit ad700c1

Browse files
committed
Add save load
1 parent 349a6ca commit ad700c1

File tree

2 files changed

+75
-0
lines changed

2 files changed

+75
-0
lines changed

docs/index.rst

+1
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ Welcome to Segmentation Models's documentation!
1717
encoders_timm
1818
losses
1919
metrics
20+
save_load
2021
insights
2122

2223

docs/save_load.rst

+74
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
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

Comments
 (0)