Add Apple's MobileOne encoder (#693)

* add mobileone encoder

* remove testing code

* add custom state dict

* remove old import

* fix formatting, extend README

* fix formatting, add limitation notice

* add support for 1 channel

* add mobileone to encoder docs

* fix encoder docs

Co-authored-by: Kevin Bondzio <[email protected]>

Author: kevinpl07

README.md

@@ -375,6 +375,25 @@ Limitations:
 </div>
 </details>
 
+<details>
+<summary style="margin-left: 25px;">MobileOne</summary>
+<div style="margin-left: 25px;">
+
+Apple's "sub-one-ms" Backbone pretrained on Imagenet! Can be used with all decoders.
+
+Note: In the official github repo the s0 variant has additional num_conv_branches, leading to more params than s1.
+
+|Encoder                         |Weights                         |Params, M                       |
+|--------------------------------|:------------------------------:|:------------------------------:|
+|mobileone_s0                    |imagenet                        |4.6M                              |
+|mobileone_s1                    |imagenet                        |4.0M                              |
+|mobileone_s2                    |imagenet                        |6.5M                              |
+|mobileone_s3                    |imagenet                        |8.8M                              |
+|mobileone_s4                    |imagenet                        |13.6M                             |
+
+</div>
+</details>
+
 
 \* `ssl`, `swsl` - semi-supervised and weakly-supervised learning on ImageNet ([repo](https://github.com/facebookresearch/semi-supervised-ImageNet1K-models)).
 

docs/encoders.rst

@@ -344,3 +344,20 @@ Mix Visual Transformer
 +-----------+----------+------------+
 | mit\_b5   | imagenet | 81M        |
 +-----------+----------+------------+
+
+MobileOne
+~~~~~~~~~~~~~~~~~~~~~
+
++-----------------+----------+------------+
+| Encoder         | Weights  | Params, M  |
++=================+==========+============+
+| mobileone\_s0   | imagenet | 4.6M       |
++-----------------+----------+------------+
+| mobileone\_s1   | imagenet | 4.0M       |
++-----------------+----------+------------+
+| mobileone\_s2   | imagenet | 6.5M       |
++-----------------+----------+------------+
+| mobileone\_s3   | imagenet | 8.8M       |
++-----------------+----------+------------+
+| mobileone\_s4   | imagenet | 13.6M      |
++-----------------+----------+------------+

segmentation_models_pytorch/encoders/__init__.py

@@ -20,6 +20,7 @@
 from .timm_mobilenetv3 import timm_mobilenetv3_encoders
 from .timm_gernet import timm_gernet_encoders
 from .mix_transformer import mix_transformer_encoders
+from .mobileone import mobileone_encoders
 
 from .timm_universal import TimmUniversalEncoder
 
@@ -44,6 +45,7 @@
 encoders.update(timm_mobilenetv3_encoders)
 encoders.update(timm_gernet_encoders)
 encoders.update(mix_transformer_encoders)
+encoders.update(mobileone_encoders)
 
 
 def get_encoder(name, in_channels=3, depth=5, weights=None, output_stride=32, **kwargs):