Skip to content

Commit b84bcb6

Browse files
committed
Update docs
1 parent c23bbb0 commit b84bcb6

File tree

1 file changed

+34
-1
lines changed

1 file changed

+34
-1
lines changed

docs/save_load.rst

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,37 @@ For example:
5959
# Or saved and pushed to the Hub simultaneously
6060
model.save_pretrained('username/my-model', push_to_hub=True, metrics={'accuracy': 0.95}, dataset='my_dataset')
6161
62+
Saving with preprocessing transform (Albumentations)
63+
----------------------------------------------------
64+
65+
You can save the preprocessing transform along with the model and push it to the Hub.
66+
This can be useful when you want to share the model with the preprocessing transform that was used during training,
67+
to make sure that the inference pipeline is consistent with the training pipeline.
68+
69+
.. code:: python
70+
71+
import albumentations as A
72+
import segmentation_models_pytorch as smp
73+
74+
# Define a preprocessing transform for image that would be used during inference
75+
preprocessing_transform = A.Compose([
76+
A.Resize(256, 256),
77+
A.Normalize()
78+
])
79+
80+
model = smp.Unet()
81+
82+
directory_or_repo_on_the_hub = "qubvel-hf/unet-with-transform" # <username>/<repo-name>
83+
84+
# Save the model and transform (and pus ot hub, if needed)
85+
model.save_pretrained(directory_or_repo_on_the_hub, push_to_hub=True)
86+
preprocessing_transform.save_pretrained(directory_or_repo_on_the_hub, push_to_hub=True)
87+
88+
# Loading transform and model
89+
restored_model = smp.from_pretrained(directory_or_repo_on_the_hub)
90+
restored_transform = A.Compose.from_pretrained(directory_or_repo_on_the_hub)
91+
92+
print(restored_transform)
6293
6394
Conclusion
6495
----------
@@ -71,4 +102,6 @@ By following these steps, you can easily save, share, and load your models, faci
71102
:target: https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/main/examples/binary_segmentation_intro.ipynb
72103
:alt: Open In Colab
73104

74-
105+
.. |colab-badge| image:: https://colab.research.google.com/assets/colab-badge.svg
106+
:target: https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/main/examples/save_load_model_and_share_with_hf_hub.ipynb
107+
:alt: Open In Colab

0 commit comments

Comments
 (0)