@@ -59,6 +59,37 @@ For example:
59
59
# Or saved and pushed to the Hub simultaneously
60
60
model.save_pretrained(' username/my-model' , push_to_hub = True , metrics = {' accuracy' : 0.95 }, dataset = ' my_dataset' )
61
61
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)
62
93
63
94
Conclusion
64
95
----------
@@ -71,4 +102,6 @@ By following these steps, you can easily save, share, and load your models, faci
71
102
:target: https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/main/examples/binary_segmentation_intro.ipynb
72
103
:alt: Open In Colab
73
104
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