segmentation_models.pytorch Documentation

import segmentation_models_pytorch as smp

model = smp.Unet(
    encoder_name="resnet34",        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
    encoder_weights="imagenet",     # use `imagenet` pre-trained weights for encoder initialization
    in_channels=1,                  # model input channels (1 for gray-scale images, 3 for RGB, etc.)
    classes=3,                      # model output channels (number of classes in your dataset)
)

from segmentation_models_pytorch.encoders import get_preprocessing_fn

preprocess_input = get_preprocessing_fn('resnet18', pretrained='imagenet')

model = smp.FPN('resnet34', in_channels=1)
mask = model(torch.ones([1, 1, 64, 64]))

aux_params=dict(
    pooling='avg',             # one of 'avg', 'max'
    dropout=0.5,               # dropout ratio, default is None
    activation='sigmoid',      # activation function, default is None
    classes=4,                 # define number of output labels
)
model = smp.Unet('resnet34', classes=4, aux_params=aux_params)
mask, label = model(x)

model = smp.Unet('resnet34', encoder_depth=4)

$ pip install segmentation-models-pytorch

$ pip install git+https://github.com/qubvel/segmentation_models.pytorch

make install_dev  # Create .venv, install SMP in dev mode

make test          # Run tests suite with pytest
make fixup         # Ruff for formatting and lint checks

make table        # Generates a table with encoders and print to stdout

@misc{Iakubovskii:2019,
  Author = {Pavel Iakubovskii},
  Title = {Segmentation Models Pytorch},
  Year = {2019},
  Publisher = {GitHub},
  Journal = {GitHub repository},
  Howpublished = {\url{https://github.com/qubvel/segmentation_models.pytorch}}
}

from segmentation_models_pytorch import Unet

# Initialize Unet with ResNet34 encoder pre-trained on ImageNet
model = Unet(encoder_name="resnet34", encoder_weights="imagenet")

To use following encoders you have to add prefix ``tu-``, e.g. ``tu-adv_inception_v3``

These models typically produce feature maps at the following downsampling scales relative to the input resolution: 1/2, 1/4, 1/8, 1/16, and 1/32

+----------------------------------+------------------+
| Encoder name                     | Support dilation |
+==================================+==================+
| bat_resnext26ts                  |        ✅        |
+----------------------------------+------------------+
| botnet26t_256                    |        ✅        |
+----------------------------------+------------------+
| botnet50ts_256                   |        ✅        |
+----------------------------------+------------------+
| coatnet_0_224                    |                  |
+----------------------------------+------------------+
| coatnet_0_rw_224                 |                  |
+----------------------------------+------------------+
| coatnet_1_224                    |                  |
+----------------------------------+------------------+
| coatnet_1_rw_224                 |                  |
+----------------------------------+------------------+
| coatnet_2_224                    |                  |
+----------------------------------+------------------+
| coatnet_2_rw_224                 |                  |
+----------------------------------+------------------+
| coatnet_3_224                    |                  |
+----------------------------------+------------------+
| coatnet_3_rw_224                 |                  |
+----------------------------------+------------------+
| coatnet_4_224                    |                  |
+----------------------------------+------------------+
| coatnet_5_224                    |                  |
+----------------------------------+------------------+
| coatnet_bn_0_rw_224              |                  |
+----------------------------------+------------------+
| coatnet_nano_cc_224              |                  |
+----------------------------------+------------------+
| coatnet_nano_rw_224              |                  |
+----------------------------------+------------------+
| coatnet_pico_rw_224              |                  |
+----------------------------------+------------------+
| coatnet_rmlp_0_rw_224            |                  |
+----------------------------------+------------------+
| coatnet_rmlp_1_rw2_224           |                  |
+----------------------------------+------------------+
| coatnet_rmlp_1_rw_224            |                  |
+----------------------------------+------------------+
| coatnet_rmlp_2_rw_224            |                  |
+----------------------------------+------------------+
| coatnet_rmlp_2_rw_384            |                  |
+----------------------------------+------------------+
| coatnet_rmlp_3_rw_224            |                  |
+----------------------------------+------------------+
| coatnet_rmlp_nano_rw_224         |                  |
+----------------------------------+------------------+
| coatnext_nano_rw_224             |                  |
+----------------------------------+------------------+
| cs3darknet_focus_l               |        ✅        |
+----------------------------------+------------------+
| cs3darknet_focus_m               |        ✅        |
+----------------------------------+------------------+
| cs3darknet_focus_s               |        ✅        |
+----------------------------------+------------------+
| cs3darknet_focus_x               |        ✅        |
+----------------------------------+------------------+
| cs3darknet_l                     |        ✅        |
+----------------------------------+------------------+
| cs3darknet_m                     |        ✅        |
+----------------------------------+------------------+
| cs3darknet_s                     |        ✅        |
+----------------------------------+------------------+
| cs3darknet_x                     |        ✅        |
+----------------------------------+------------------+
| cs3edgenet_x                     |        ✅        |
+----------------------------------+------------------+
| cs3se_edgenet_x                  |        ✅        |
+----------------------------------+------------------+
| cs3sedarknet_l                   |        ✅        |
+----------------------------------+------------------+
| cs3sedarknet_x                   |        ✅        |
+----------------------------------+------------------+
| cs3sedarknet_xdw                 |        ✅        |
+----------------------------------+------------------+
| cspdarknet53                     |        ✅        |
+----------------------------------+------------------+
| cspresnet50                      |        ✅        |
+----------------------------------+------------------+
| cspresnet50d                     |        ✅        |
+----------------------------------+------------------+
| cspresnet50w                     |        ✅        |
+----------------------------------+------------------+
| cspresnext50                     |        ✅        |
+----------------------------------+------------------+
| darknet17                        |        ✅        |
+----------------------------------+------------------+
| darknet21                        |        ✅        |
+----------------------------------+------------------+
| darknet53                        |        ✅        |
+----------------------------------+------------------+
| darknetaa53                      |        ✅        |
+----------------------------------+------------------+
| densenet121                      |                  |
+----------------------------------+------------------+
| densenet161                      |                  |
+----------------------------------+------------------+
| densenet169                      |                  |
+----------------------------------+------------------+
| densenet201                      |                  |
+----------------------------------+------------------+
| densenet264d                     |                  |
+----------------------------------+------------------+
| densenetblur121d                 |                  |
+----------------------------------+------------------+
| dla34                            |                  |
+----------------------------------+------------------+
| dla46_c                          |                  |
+----------------------------------+------------------+
| dla46x_c                         |                  |
+----------------------------------+------------------+
| dla60                            |                  |
+----------------------------------+------------------+
| dla60_res2net                    |                  |
+----------------------------------+------------------+
| dla60_res2next                   |                  |
+----------------------------------+------------------+
| dla60x                           |                  |
+----------------------------------+------------------+
| dla60x_c                         |                  |
+----------------------------------+------------------+
| dla102                           |                  |
+----------------------------------+------------------+
| dla102x                          |                  |
+----------------------------------+------------------+
| dla102x2                         |                  |
+----------------------------------+------------------+
| dla169                           |                  |
+----------------------------------+------------------+
| dm_nfnet_f0                      |        ✅        |
+----------------------------------+------------------+
| dm_nfnet_f1                      |        ✅        |
+----------------------------------+------------------+
| dm_nfnet_f2                      |        ✅        |
+----------------------------------+------------------+
| dm_nfnet_f3                      |        ✅        |
+----------------------------------+------------------+
| dm_nfnet_f4                      |        ✅        |
+----------------------------------+------------------+
| dm_nfnet_f5                      |        ✅        |
+----------------------------------+------------------+
| dm_nfnet_f6                      |        ✅        |
+----------------------------------+------------------+
| dpn48b                           |                  |
+----------------------------------+------------------+
| dpn68                            |                  |
+----------------------------------+------------------+
| dpn68b                           |                  |
+----------------------------------+------------------+
| dpn92                            |                  |
+----------------------------------+------------------+
| dpn98                            |                  |
+----------------------------------+------------------+
| dpn107                           |                  |
+----------------------------------+------------------+
| dpn131                           |                  |
+----------------------------------+------------------+
| eca_botnext26ts_256              |        ✅        |
+----------------------------------+------------------+
| eca_halonext26ts                 |        ✅        |
+----------------------------------+------------------+
| eca_nfnet_l0                     |        ✅        |
+----------------------------------+------------------+
| eca_nfnet_l1                     |        ✅        |
+----------------------------------+------------------+
| eca_nfnet_l2                     |        ✅        |
+----------------------------------+------------------+
| eca_nfnet_l3                     |        ✅        |
+----------------------------------+------------------+
| eca_resnet33ts                   |        ✅        |
+----------------------------------+------------------+
| eca_resnext26ts                  |        ✅        |
+----------------------------------+------------------+
| eca_vovnet39b                    |                  |
+----------------------------------+------------------+
| ecaresnet26t                     |        ✅        |
+----------------------------------+------------------+
| ecaresnet50d                     |        ✅        |
+----------------------------------+------------------+
| ecaresnet50d_pruned              |        ✅        |
+----------------------------------+------------------+
| ecaresnet50t                     |        ✅        |
+----------------------------------+------------------+
| ecaresnet101d                    |        ✅        |
+----------------------------------+------------------+
| ecaresnet101d_pruned             |        ✅        |
+----------------------------------+------------------+
| ecaresnet200d                    |        ✅        |
+----------------------------------+------------------+
| ecaresnet269d                    |        ✅        |
+----------------------------------+------------------+
| ecaresnetlight                   |        ✅        |
+----------------------------------+------------------+
| ecaresnext26t_32x4d              |        ✅        |
+----------------------------------+------------------+
| ecaresnext50t_32x4d              |        ✅        |
+----------------------------------+------------------+
| efficientnet_b0                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_b0_g8_gn            |        ✅        |
+----------------------------------+------------------+
| efficientnet_b0_g16_evos         |        ✅        |
+----------------------------------+------------------+
| efficientnet_b0_gn               |        ✅        |
+----------------------------------+------------------+
| efficientnet_b1                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_b1_pruned           |        ✅        |
+----------------------------------+------------------+
| efficientnet_b2                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_b2_pruned           |        ✅        |
+----------------------------------+------------------+
| efficientnet_b3                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_b3_g8_gn            |        ✅        |
+----------------------------------+------------------+
| efficientnet_b3_gn               |        ✅        |
+----------------------------------+------------------+
| efficientnet_b3_pruned           |        ✅        |
+----------------------------------+------------------+
| efficientnet_b4                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_b5                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_b6                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_b7                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_b8                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_blur_b0             |        ✅        |
+----------------------------------+------------------+
| efficientnet_cc_b0_4e            |        ✅        |
+----------------------------------+------------------+
| efficientnet_cc_b0_8e            |        ✅        |
+----------------------------------+------------------+
| efficientnet_cc_b1_8e            |        ✅        |
+----------------------------------+------------------+
| efficientnet_el                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_el_pruned           |        ✅        |
+----------------------------------+------------------+
| efficientnet_em                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_es                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_es_pruned           |        ✅        |
+----------------------------------+------------------+
| efficientnet_l2                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_lite0               |        ✅        |
+----------------------------------+------------------+
| efficientnet_lite1               |        ✅        |
+----------------------------------+------------------+
| efficientnet_lite2               |        ✅        |
+----------------------------------+------------------+
| efficientnet_lite3               |        ✅        |
+----------------------------------+------------------+
| efficientnet_lite4               |        ✅        |
+----------------------------------+------------------+
| efficientnetv2_l                 |        ✅        |
+----------------------------------+------------------+
| efficientnetv2_m                 |        ✅        |
+----------------------------------+------------------+
| efficientnetv2_rw_m              |        ✅        |
+----------------------------------+------------------+
| efficientnetv2_rw_s              |        ✅        |
+----------------------------------+------------------+
| efficientnetv2_rw_t              |        ✅        |
+----------------------------------+------------------+
| efficientnetv2_s                 |        ✅        |
+----------------------------------+------------------+
| efficientnetv2_xl                |        ✅        |
+----------------------------------+------------------+
| ese_vovnet19b_dw                 |                  |
+----------------------------------+------------------+
| ese_vovnet19b_slim               |                  |
+----------------------------------+------------------+
| ese_vovnet19b_slim_dw            |                  |
+----------------------------------+------------------+
| ese_vovnet39b                    |                  |
+----------------------------------+------------------+
| ese_vovnet39b_evos               |                  |
+----------------------------------+------------------+
| ese_vovnet57b                    |                  |
+----------------------------------+------------------+
| ese_vovnet99b                    |                  |
+----------------------------------+------------------+
| fbnetc_100                       |        ✅        |
+----------------------------------+------------------+
| fbnetv3_b                        |        ✅        |
+----------------------------------+------------------+
| fbnetv3_d                        |        ✅        |
+----------------------------------+------------------+
| fbnetv3_g                        |        ✅        |
+----------------------------------+------------------+
| gc_efficientnetv2_rw_t           |        ✅        |
+----------------------------------+------------------+
| gcresnet33ts                     |        ✅        |
+----------------------------------+------------------+
| gcresnet50t                      |        ✅        |
+----------------------------------+------------------+
| gcresnext26ts                    |        ✅        |
+----------------------------------+------------------+
| gcresnext50ts                    |        ✅        |
+----------------------------------+------------------+
| gernet_l                         |        ✅        |
+----------------------------------+------------------+
| gernet_m                         |        ✅        |
+----------------------------------+------------------+
| gernet_s                         |        ✅        |
+----------------------------------+------------------+
| ghostnet_050                     |                  |
+----------------------------------+------------------+
| ghostnet_100                     |                  |
+----------------------------------+------------------+
| ghostnet_130                     |                  |
+----------------------------------+------------------+
| ghostnetv2_100                   |                  |
+----------------------------------+------------------+
| ghostnetv2_130                   |                  |
+----------------------------------+------------------+
| ghostnetv2_160                   |                  |
+----------------------------------+------------------+
| halo2botnet50ts_256              |        ✅        |
+----------------------------------+------------------+
| halonet26t                       |        ✅        |
+----------------------------------+------------------+
| halonet50ts                      |        ✅        |
+----------------------------------+------------------+
| halonet_h1                       |        ✅        |
+----------------------------------+------------------+
| haloregnetz_b                    |        ✅        |
+----------------------------------+------------------+
| hardcorenas_a                    |        ✅        |
+----------------------------------+------------------+
| hardcorenas_b                    |        ✅        |
+----------------------------------+------------------+
| hardcorenas_c                    |        ✅        |
+----------------------------------+------------------+
| hardcorenas_d                    |        ✅        |
+----------------------------------+------------------+
| hardcorenas_e                    |        ✅        |
+----------------------------------+------------------+
| hardcorenas_f                    |        ✅        |
+----------------------------------+------------------+
| hrnet_w18                        |                  |
+----------------------------------+------------------+
| hrnet_w18_small                  |                  |
+----------------------------------+------------------+
| hrnet_w18_small_v2               |                  |
+----------------------------------+------------------+
| hrnet_w18_ssld                   |                  |
+----------------------------------+------------------+
| hrnet_w30                        |                  |
+----------------------------------+------------------+
| hrnet_w32                        |                  |
+----------------------------------+------------------+
| hrnet_w40                        |                  |
+----------------------------------+------------------+
| hrnet_w44                        |                  |
+----------------------------------+------------------+
| hrnet_w48                        |                  |
+----------------------------------+------------------+
| hrnet_w48_ssld                   |                  |
+----------------------------------+------------------+
| hrnet_w64                        |                  |
+----------------------------------+------------------+
| lambda_resnet26rpt_256           |        ✅        |
+----------------------------------+------------------+
| lambda_resnet26t                 |        ✅        |
+----------------------------------+------------------+
| lambda_resnet50ts                |        ✅        |
+----------------------------------+------------------+
| lamhalobotnet50ts_256            |        ✅        |
+----------------------------------+------------------+
| lcnet_035                        |        ✅        |
+----------------------------------+------------------+
| lcnet_050                        |        ✅        |
+----------------------------------+------------------+
| lcnet_075                        |        ✅        |
+----------------------------------+------------------+
| lcnet_100                        |        ✅        |
+----------------------------------+------------------+
| lcnet_150                        |        ✅        |
+----------------------------------+------------------+
| legacy_senet154                  |                  |
+----------------------------------+------------------+
| legacy_seresnet18                |                  |
+----------------------------------+------------------+
| legacy_seresnet34                |                  |
+----------------------------------+------------------+
| legacy_seresnet50                |                  |
+----------------------------------+------------------+
| legacy_seresnet101               |                  |
+----------------------------------+------------------+
| legacy_seresnet152               |                  |
+----------------------------------+------------------+
| legacy_seresnext26_32x4d         |                  |
+----------------------------------+------------------+
| legacy_seresnext50_32x4d         |                  |
+----------------------------------+------------------+
| legacy_seresnext101_32x4d        |                  |
+----------------------------------+------------------+
| maxvit_base_tf_224               |                  |
+----------------------------------+------------------+
| maxvit_base_tf_384               |                  |
+----------------------------------+------------------+
| maxvit_base_tf_512               |                  |
+----------------------------------+------------------+
| maxvit_large_tf_224              |                  |
+----------------------------------+------------------+
| maxvit_large_tf_384              |                  |
+----------------------------------+------------------+
| maxvit_large_tf_512              |                  |
+----------------------------------+------------------+
| maxvit_nano_rw_256               |                  |
+----------------------------------+------------------+
| maxvit_pico_rw_256               |                  |
+----------------------------------+------------------+
| maxvit_rmlp_base_rw_224          |                  |
+----------------------------------+------------------+
| maxvit_rmlp_base_rw_384          |                  |
+----------------------------------+------------------+
| maxvit_rmlp_nano_rw_256          |                  |
+----------------------------------+------------------+
| maxvit_rmlp_pico_rw_256          |                  |
+----------------------------------+------------------+
| maxvit_rmlp_small_rw_224         |                  |
+----------------------------------+------------------+
| maxvit_rmlp_small_rw_256         |                  |
+----------------------------------+------------------+
| maxvit_rmlp_tiny_rw_256          |                  |
+----------------------------------+------------------+
| maxvit_small_tf_224              |                  |
+----------------------------------+------------------+
| maxvit_small_tf_384              |                  |
+----------------------------------+------------------+
| maxvit_small_tf_512              |                  |
+----------------------------------+------------------+
| maxvit_tiny_pm_256               |                  |
+----------------------------------+------------------+
| maxvit_tiny_rw_224               |                  |
+----------------------------------+------------------+
| maxvit_tiny_rw_256               |                  |
+----------------------------------+------------------+
| maxvit_tiny_tf_224               |                  |
+----------------------------------+------------------+
| maxvit_tiny_tf_384               |                  |
+----------------------------------+------------------+
| maxvit_tiny_tf_512               |                  |
+----------------------------------+------------------+
| maxvit_xlarge_tf_224             |                  |
+----------------------------------+------------------+
| maxvit_xlarge_tf_384             |                  |
+----------------------------------+------------------+
| maxvit_xlarge_tf_512             |                  |
+----------------------------------+------------------+
| maxxvit_rmlp_nano_rw_256         |                  |
+----------------------------------+------------------+
| maxxvit_rmlp_small_rw_256        |                  |
+----------------------------------+------------------+
| maxxvit_rmlp_tiny_rw_256         |                  |
+----------------------------------+------------------+
| maxxvitv2_nano_rw_256            |                  |
+----------------------------------+------------------+
| maxxvitv2_rmlp_base_rw_224       |                  |
+----------------------------------+------------------+
| maxxvitv2_rmlp_base_rw_384       |                  |
+----------------------------------+------------------+
| maxxvitv2_rmlp_large_rw_224      |                  |
+----------------------------------+------------------+
| mixnet_l                         |        ✅        |
+----------------------------------+------------------+
| mixnet_m                         |        ✅        |
+----------------------------------+------------------+
| mixnet_s                         |        ✅        |
+----------------------------------+------------------+
| mixnet_xl                        |        ✅        |
+----------------------------------+------------------+
| mixnet_xxl                       |        ✅        |
+----------------------------------+------------------+
| mnasnet_050                      |        ✅        |
+----------------------------------+------------------+
| mnasnet_075                      |        ✅        |
+----------------------------------+------------------+
| mnasnet_100                      |        ✅        |
+----------------------------------+------------------+
| mnasnet_140                      |        ✅        |
+----------------------------------+------------------+
| mnasnet_small                    |        ✅        |
+----------------------------------+------------------+
| mobilenet_edgetpu_100            |        ✅        |
+----------------------------------+------------------+
| mobilenet_edgetpu_v2_l           |        ✅        |
+----------------------------------+------------------+
| mobilenet_edgetpu_v2_m           |        ✅        |
+----------------------------------+------------------+
| mobilenet_edgetpu_v2_s           |        ✅        |
+----------------------------------+------------------+
| mobilenet_edgetpu_v2_xs          |        ✅        |
+----------------------------------+------------------+
| mobilenetv1_100                  |        ✅        |
+----------------------------------+------------------+
| mobilenetv1_100h                 |        ✅        |
+----------------------------------+------------------+
| mobilenetv1_125                  |        ✅        |
+----------------------------------+------------------+
| mobilenetv2_035                  |        ✅        |
+----------------------------------+------------------+
| mobilenetv2_050                  |        ✅        |
+----------------------------------+------------------+
| mobilenetv2_075                  |        ✅        |
+----------------------------------+------------------+
| mobilenetv2_100                  |        ✅        |
+----------------------------------+------------------+
| mobilenetv2_110d                 |        ✅        |
+----------------------------------+------------------+
| mobilenetv2_120d                 |        ✅        |
+----------------------------------+------------------+
| mobilenetv2_140                  |        ✅        |
+----------------------------------+------------------+
| mobilenetv3_large_075            |        ✅        |
+----------------------------------+------------------+
| mobilenetv3_large_100            |        ✅        |
+----------------------------------+------------------+
| mobilenetv3_large_150d           |        ✅        |
+----------------------------------+------------------+
| mobilenetv3_rw                   |        ✅        |
+----------------------------------+------------------+
| mobilenetv3_small_050            |        ✅        |
+----------------------------------+------------------+
| mobilenetv3_small_075            |        ✅        |
+----------------------------------+------------------+
| mobilenetv3_small_100            |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_conv_aa_large        |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_conv_aa_medium       |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_conv_blur_medium     |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_conv_large           |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_conv_medium          |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_conv_small           |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_conv_small_035       |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_conv_small_050       |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_hybrid_large         |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_hybrid_large_075     |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_hybrid_medium        |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_hybrid_medium_075    |        ✅        |
+----------------------------------+------------------+
| mobileone_s0                     |        ✅        |
+----------------------------------+------------------+
| mobileone_s1                     |        ✅        |
+----------------------------------+------------------+
| mobileone_s2                     |        ✅        |
+----------------------------------+------------------+
| mobileone_s3                     |        ✅        |
+----------------------------------+------------------+
| mobileone_s4                     |        ✅        |
+----------------------------------+------------------+
| mobilevit_s                      |        ✅        |
+----------------------------------+------------------+
| mobilevit_xs                     |        ✅        |
+----------------------------------+------------------+
| mobilevit_xxs                    |        ✅        |
+----------------------------------+------------------+
| mobilevitv2_050                  |        ✅        |
+----------------------------------+------------------+
| mobilevitv2_075                  |        ✅        |
+----------------------------------+------------------+
| mobilevitv2_100                  |        ✅        |
+----------------------------------+------------------+
| mobilevitv2_125                  |        ✅        |
+----------------------------------+------------------+
| mobilevitv2_150                  |        ✅        |
+----------------------------------+------------------+
| mobilevitv2_175                  |        ✅        |
+----------------------------------+------------------+
| mobilevitv2_200                  |        ✅        |
+----------------------------------+------------------+
| nf_ecaresnet26                   |        ✅        |
+----------------------------------+------------------+
| nf_ecaresnet50                   |        ✅        |
+----------------------------------+------------------+
| nf_ecaresnet101                  |        ✅        |
+----------------------------------+------------------+
| nf_regnet_b0                     |        ✅        |
+----------------------------------+------------------+
| nf_regnet_b1                     |        ✅        |
+----------------------------------+------------------+
| nf_regnet_b2                     |        ✅        |
+----------------------------------+------------------+
| nf_regnet_b3                     |        ✅        |
+----------------------------------+------------------+
| nf_regnet_b4                     |        ✅        |
+----------------------------------+------------------+
| nf_regnet_b5                     |        ✅        |
+----------------------------------+------------------+
| nf_resnet26                      |        ✅        |
+----------------------------------+------------------+
| nf_resnet50                      |        ✅        |
+----------------------------------+------------------+
| nf_resnet101                     |        ✅        |
+----------------------------------+------------------+
| nf_seresnet26                    |        ✅        |
+----------------------------------+------------------+
| nf_seresnet50                    |        ✅        |
+----------------------------------+------------------+
| nf_seresnet101                   |        ✅        |
+----------------------------------+------------------+
| nfnet_f0                         |        ✅        |
+----------------------------------+------------------+
| nfnet_f1                         |        ✅        |
+----------------------------------+------------------+
| nfnet_f2                         |        ✅        |
+----------------------------------+------------------+
| nfnet_f3                         |        ✅        |
+----------------------------------+------------------+
| nfnet_f4                         |        ✅        |
+----------------------------------+------------------+
| nfnet_f5                         |        ✅        |
+----------------------------------+------------------+
| nfnet_f6                         |        ✅        |
+----------------------------------+------------------+
| nfnet_f7                         |        ✅        |
+----------------------------------+------------------+
| nfnet_l0                         |        ✅        |
+----------------------------------+------------------+
| regnetv_040                      |        ✅        |
+----------------------------------+------------------+
| regnetv_064                      |        ✅        |
+----------------------------------+------------------+
| regnetx_002                      |        ✅        |
+----------------------------------+------------------+
| regnetx_004                      |        ✅        |
+----------------------------------+------------------+
| regnetx_004_tv                   |        ✅        |
+----------------------------------+------------------+
| regnetx_006                      |        ✅        |
+----------------------------------+------------------+
| regnetx_008                      |        ✅        |
+----------------------------------+------------------+
| regnetx_016                      |        ✅        |
+----------------------------------+------------------+
| regnetx_032                      |        ✅        |
+----------------------------------+------------------+
| regnetx_040                      |        ✅        |
+----------------------------------+------------------+
| regnetx_064                      |        ✅        |
+----------------------------------+------------------+
| regnetx_080                      |        ✅        |
+----------------------------------+------------------+
| regnetx_120                      |        ✅        |
+----------------------------------+------------------+
| regnetx_160                      |        ✅        |
+----------------------------------+------------------+
| regnetx_320                      |        ✅        |
+----------------------------------+------------------+
| regnety_002                      |        ✅        |
+----------------------------------+------------------+
| regnety_004                      |        ✅        |
+----------------------------------+------------------+
| regnety_006                      |        ✅        |
+----------------------------------+------------------+
| regnety_008                      |        ✅        |
+----------------------------------+------------------+
| regnety_008_tv                   |        ✅        |
+----------------------------------+------------------+
| regnety_016                      |        ✅        |
+----------------------------------+------------------+
| regnety_032                      |        ✅        |
+----------------------------------+------------------+
| regnety_040                      |        ✅        |
+----------------------------------+------------------+
| regnety_040_sgn                  |        ✅        |
+----------------------------------+------------------+
| regnety_064                      |        ✅        |
+----------------------------------+------------------+
| regnety_080                      |        ✅        |
+----------------------------------+------------------+
| regnety_080_tv                   |        ✅        |
+----------------------------------+------------------+
| regnety_120                      |        ✅        |
+----------------------------------+------------------+
| regnety_160                      |        ✅        |
+----------------------------------+------------------+
| regnety_320                      |        ✅        |
+----------------------------------+------------------+
| regnety_640                      |        ✅        |
+----------------------------------+------------------+
| regnety_1280                     |        ✅        |
+----------------------------------+------------------+
| regnety_2560                     |        ✅        |
+----------------------------------+------------------+
| regnetz_005                      |        ✅        |
+----------------------------------+------------------+
| regnetz_040                      |        ✅        |
+----------------------------------+------------------+
| regnetz_040_h                    |        ✅        |
+----------------------------------+------------------+
| regnetz_b16                      |        ✅        |
+----------------------------------+------------------+
| regnetz_b16_evos                 |        ✅        |
+----------------------------------+------------------+
| regnetz_c16                      |        ✅        |
+----------------------------------+------------------+
| regnetz_c16_evos                 |        ✅        |
+----------------------------------+------------------+
| regnetz_d8                       |        ✅        |
+----------------------------------+------------------+
| regnetz_d8_evos                  |        ✅        |
+----------------------------------+------------------+
| regnetz_d32                      |        ✅        |
+----------------------------------+------------------+
| regnetz_e8                       |        ✅        |
+----------------------------------+------------------+
| repghostnet_050                  |                  |
+----------------------------------+------------------+
| repghostnet_058                  |                  |
+----------------------------------+------------------+
| repghostnet_080                  |                  |
+----------------------------------+------------------+
| repghostnet_100                  |                  |
+----------------------------------+------------------+
| repghostnet_111                  |                  |
+----------------------------------+------------------+
| repghostnet_130                  |                  |
+----------------------------------+------------------+
| repghostnet_150                  |                  |
+----------------------------------+------------------+
| repghostnet_200                  |                  |
+----------------------------------+------------------+
| repvgg_a0                        |        ✅        |
+----------------------------------+------------------+
| repvgg_a1                        |        ✅        |
+----------------------------------+------------------+
| repvgg_a2                        |        ✅        |
+----------------------------------+------------------+
| repvgg_b0                        |        ✅        |
+----------------------------------+------------------+
| repvgg_b1                        |        ✅        |
+----------------------------------+------------------+
| repvgg_b1g4                      |        ✅        |
+----------------------------------+------------------+
| repvgg_b2                        |        ✅        |
+----------------------------------+------------------+
| repvgg_b2g4                      |        ✅        |
+----------------------------------+------------------+
| repvgg_b3                        |        ✅        |
+----------------------------------+------------------+
| repvgg_b3g4                      |        ✅        |
+----------------------------------+------------------+
| repvgg_d2se                      |        ✅        |
+----------------------------------+------------------+
| res2net50_14w_8s                 |        ✅        |
+----------------------------------+------------------+
| res2net50_26w_4s                 |        ✅        |
+----------------------------------+------------------+
| res2net50_26w_6s                 |        ✅        |
+----------------------------------+------------------+
| res2net50_26w_8s                 |        ✅        |
+----------------------------------+------------------+
| res2net50_48w_2s                 |        ✅        |
+----------------------------------+------------------+
| res2net50d                       |        ✅        |
+----------------------------------+------------------+
| res2net101_26w_4s                |        ✅        |
+----------------------------------+------------------+
| res2net101d                      |        ✅        |
+----------------------------------+------------------+
| res2next50                       |        ✅        |
+----------------------------------+------------------+
| resnest14d                       |        ✅        |
+----------------------------------+------------------+
| resnest26d                       |        ✅        |
+----------------------------------+------------------+
| resnest50d                       |        ✅        |
+----------------------------------+------------------+
| resnest50d_1s4x24d               |        ✅        |
+----------------------------------+------------------+
| resnest50d_4s2x40d               |        ✅        |
+----------------------------------+------------------+
| resnest101e                      |        ✅        |
+----------------------------------+------------------+
| resnest200e                      |        ✅        |
+----------------------------------+------------------+
| resnest269e                      |        ✅        |
+----------------------------------+------------------+
| resnet10t                        |        ✅        |
+----------------------------------+------------------+
| resnet14t                        |        ✅        |
+----------------------------------+------------------+
| resnet18                         |        ✅        |
+----------------------------------+------------------+
| resnet18d                        |        ✅        |
+----------------------------------+------------------+
| resnet26                         |        ✅        |
+----------------------------------+------------------+
| resnet26d                        |        ✅        |
+----------------------------------+------------------+
| resnet26t                        |        ✅        |
+----------------------------------+------------------+
| resnet32ts                       |        ✅        |
+----------------------------------+------------------+
| resnet33ts                       |        ✅        |
+----------------------------------+------------------+
| resnet34                         |        ✅        |
+----------------------------------+------------------+
| resnet34d                        |        ✅        |
+----------------------------------+------------------+
| resnet50                         |        ✅        |
+----------------------------------+------------------+
| resnet50_clip                    |        ✅        |
+----------------------------------+------------------+
| resnet50_clip_gap                |        ✅        |
+----------------------------------+------------------+
| resnet50_gn                      |        ✅        |
+----------------------------------+------------------+
| resnet50_mlp                     |        ✅        |
+----------------------------------+------------------+
| resnet50c                        |        ✅        |
+----------------------------------+------------------+
| resnet50d                        |        ✅        |
+----------------------------------+------------------+
| resnet50s                        |        ✅        |
+----------------------------------+------------------+
| resnet50t                        |        ✅        |
+----------------------------------+------------------+
| resnet50x4_clip                  |        ✅        |
+----------------------------------+------------------+
| resnet50x4_clip_gap              |        ✅        |
+----------------------------------+------------------+
| resnet50x16_clip                 |        ✅        |
+----------------------------------+------------------+
| resnet50x16_clip_gap             |        ✅        |
+----------------------------------+------------------+
| resnet50x64_clip                 |        ✅        |
+----------------------------------+------------------+
| resnet50x64_clip_gap             |        ✅        |
+----------------------------------+------------------+
| resnet51q                        |        ✅        |
+----------------------------------+------------------+
| resnet61q                        |        ✅        |
+----------------------------------+------------------+
| resnet101                        |        ✅        |
+----------------------------------+------------------+
| resnet101_clip                   |        ✅        |
+----------------------------------+------------------+
| resnet101_clip_gap               |        ✅        |
+----------------------------------+------------------+
| resnet101c                       |        ✅        |
+----------------------------------+------------------+
| resnet101d                       |        ✅        |
+----------------------------------+------------------+
| resnet101s                       |        ✅        |
+----------------------------------+------------------+
| resnet152                        |        ✅        |
+----------------------------------+------------------+
| resnet152c                       |        ✅        |
+----------------------------------+------------------+
| resnet152d                       |        ✅        |
+----------------------------------+------------------+
| resnet152s                       |        ✅        |
+----------------------------------+------------------+
| resnet200                        |        ✅        |
+----------------------------------+------------------+
| resnet200d                       |        ✅        |
+----------------------------------+------------------+
| resnetaa34d                      |        ✅        |
+----------------------------------+------------------+
| resnetaa50                       |        ✅        |
+----------------------------------+------------------+
| resnetaa50d                      |        ✅        |
+----------------------------------+------------------+
| resnetaa101d                     |        ✅        |
+----------------------------------+------------------+
| resnetblur18                     |        ✅        |
+----------------------------------+------------------+
| resnetblur50                     |        ✅        |
+----------------------------------+------------------+
| resnetblur50d                    |        ✅        |
+----------------------------------+------------------+
| resnetblur101d                   |        ✅        |
+----------------------------------+------------------+
| resnetrs50                       |        ✅        |
+----------------------------------+------------------+
| resnetrs101                      |        ✅        |
+----------------------------------+------------------+
| resnetrs152                      |        ✅        |
+----------------------------------+------------------+
| resnetrs200                      |        ✅        |
+----------------------------------+------------------+
| resnetrs270                      |        ✅        |
+----------------------------------+------------------+
| resnetrs350                      |        ✅        |
+----------------------------------+------------------+
| resnetrs420                      |        ✅        |
+----------------------------------+------------------+
| resnetv2_18                      |        ✅        |
+----------------------------------+------------------+
| resnetv2_18d                     |        ✅        |
+----------------------------------+------------------+
| resnetv2_34                      |        ✅        |
+----------------------------------+------------------+
| resnetv2_34d                     |        ✅        |
+----------------------------------+------------------+
| resnetv2_50                      |        ✅        |
+----------------------------------+------------------+
| resnetv2_50d                     |        ✅        |
+----------------------------------+------------------+
| resnetv2_50d_evos                |        ✅        |
+----------------------------------+------------------+
| resnetv2_50d_frn                 |        ✅        |
+----------------------------------+------------------+
| resnetv2_50d_gn                  |        ✅        |
+----------------------------------+------------------+
| resnetv2_50t                     |        ✅        |
+----------------------------------+------------------+
| resnetv2_50x1_bit                |        ✅        |
+----------------------------------+------------------+
| resnetv2_50x3_bit                |        ✅        |
+----------------------------------+------------------+
| resnetv2_101                     |        ✅        |
+----------------------------------+------------------+
| resnetv2_101d                    |        ✅        |
+----------------------------------+------------------+
| resnetv2_101x1_bit               |        ✅        |
+----------------------------------+------------------+
| resnetv2_101x3_bit               |        ✅        |
+----------------------------------+------------------+
| resnetv2_152                     |        ✅        |
+----------------------------------+------------------+
| resnetv2_152d                    |        ✅        |
+----------------------------------+------------------+
| resnetv2_152x2_bit               |        ✅        |
+----------------------------------+------------------+
| resnetv2_152x4_bit               |        ✅        |
+----------------------------------+------------------+
| resnext26ts                      |        ✅        |
+----------------------------------+------------------+
| resnext50_32x4d                  |        ✅        |
+----------------------------------+------------------+
| resnext50d_32x4d                 |        ✅        |
+----------------------------------+------------------+
| resnext101_32x4d                 |        ✅        |
+----------------------------------+------------------+
| resnext101_32x8d                 |        ✅        |
+----------------------------------+------------------+
| resnext101_32x16d                |        ✅        |
+----------------------------------+------------------+
| resnext101_32x32d                |        ✅        |
+----------------------------------+------------------+
| resnext101_64x4d                 |        ✅        |
+----------------------------------+------------------+
| rexnet_100                       |        ✅        |
+----------------------------------+------------------+
| rexnet_130                       |        ✅        |
+----------------------------------+------------------+
| rexnet_150                       |        ✅        |
+----------------------------------+------------------+
| rexnet_200                       |        ✅        |
+----------------------------------+------------------+
| rexnet_300                       |        ✅        |
+----------------------------------+------------------+
| rexnetr_100                      |        ✅        |
+----------------------------------+------------------+
| rexnetr_130                      |        ✅        |
+----------------------------------+------------------+
| rexnetr_150                      |        ✅        |
+----------------------------------+------------------+
| rexnetr_200                      |        ✅        |
+----------------------------------+------------------+
| rexnetr_300                      |        ✅        |
+----------------------------------+------------------+
| sebotnet33ts_256                 |        ✅        |
+----------------------------------+------------------+
| sedarknet21                      |        ✅        |
+----------------------------------+------------------+
| sehalonet33ts                    |        ✅        |
+----------------------------------+------------------+
| selecsls42                       |                  |
+----------------------------------+------------------+
| selecsls42b                      |                  |
+----------------------------------+------------------+
| selecsls60                       |                  |
+----------------------------------+------------------+
| selecsls60b                      |                  |
+----------------------------------+------------------+
| selecsls84                       |                  |
+----------------------------------+------------------+
| semnasnet_050                    |        ✅        |
+----------------------------------+------------------+
| semnasnet_075                    |        ✅        |
+----------------------------------+------------------+
| semnasnet_100                    |        ✅        |
+----------------------------------+------------------+
| semnasnet_140                    |        ✅        |
+----------------------------------+------------------+
| senet154                         |        ✅        |
+----------------------------------+------------------+
| seresnet18                       |        ✅        |
+----------------------------------+------------------+
| seresnet33ts                     |        ✅        |
+----------------------------------+------------------+
| seresnet34                       |        ✅        |
+----------------------------------+------------------+
| seresnet50                       |        ✅        |
+----------------------------------+------------------+
| seresnet50t                      |        ✅        |
+----------------------------------+------------------+
| seresnet101                      |        ✅        |
+----------------------------------+------------------+
| seresnet152                      |        ✅        |
+----------------------------------+------------------+
| seresnet152d                     |        ✅        |
+----------------------------------+------------------+
| seresnet200d                     |        ✅        |
+----------------------------------+------------------+
| seresnet269d                     |        ✅        |
+----------------------------------+------------------+
| seresnetaa50d                    |        ✅        |
+----------------------------------+------------------+
| seresnext26d_32x4d               |        ✅        |
+----------------------------------+------------------+
| seresnext26t_32x4d               |        ✅        |
+----------------------------------+------------------+
| seresnext26ts                    |        ✅        |
+----------------------------------+------------------+
| seresnext50_32x4d                |        ✅        |
+----------------------------------+------------------+
| seresnext101_32x4d               |        ✅        |
+----------------------------------+------------------+
| seresnext101_32x8d               |        ✅        |
+----------------------------------+------------------+
| seresnext101_64x4d               |        ✅        |
+----------------------------------+------------------+
| seresnext101d_32x8d              |        ✅        |
+----------------------------------+------------------+
| seresnextaa101d_32x8d            |        ✅        |
+----------------------------------+------------------+
| seresnextaa201d_32x8d            |        ✅        |
+----------------------------------+------------------+
| skresnet18                       |        ✅        |
+----------------------------------+------------------+
| skresnet34                       |        ✅        |
+----------------------------------+------------------+
| skresnet50                       |        ✅        |
+----------------------------------+------------------+
| skresnet50d                      |        ✅        |
+----------------------------------+------------------+
| skresnext50_32x4d                |        ✅        |
+----------------------------------+------------------+
| spnasnet_100                     |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_b0               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_b1               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_b2               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_b3               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_b4               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_b5               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_b6               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_b7               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_b8               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_cc_b0_4e         |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_cc_b0_8e         |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_cc_b1_8e         |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_el               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_em               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_es               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_l2               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_lite0            |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_lite1            |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_lite2            |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_lite3            |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_lite4            |        ✅        |
+----------------------------------+------------------+
| tf_efficientnetv2_b0             |        ✅        |
+----------------------------------+------------------+
| tf_efficientnetv2_b1             |        ✅        |
+----------------------------------+------------------+
| tf_efficientnetv2_b2             |        ✅        |
+----------------------------------+------------------+
| tf_efficientnetv2_b3             |        ✅        |
+----------------------------------+------------------+
| tf_efficientnetv2_l              |        ✅        |
+----------------------------------+------------------+
| tf_efficientnetv2_m              |        ✅        |
+----------------------------------+------------------+
| tf_efficientnetv2_s              |        ✅        |
+----------------------------------+------------------+
| tf_efficientnetv2_xl             |        ✅        |
+----------------------------------+------------------+
| tf_mixnet_l                      |        ✅        |
+----------------------------------+------------------+
| tf_mixnet_m                      |        ✅        |
+----------------------------------+------------------+
| tf_mixnet_s                      |        ✅        |
+----------------------------------+------------------+
| tf_mobilenetv3_large_075         |        ✅        |
+----------------------------------+------------------+
| tf_mobilenetv3_large_100         |        ✅        |
+----------------------------------+------------------+
| tf_mobilenetv3_large_minimal_100 |        ✅        |
+----------------------------------+------------------+
| tf_mobilenetv3_small_075         |        ✅        |
+----------------------------------+------------------+
| tf_mobilenetv3_small_100         |        ✅        |
+----------------------------------+------------------+
| tf_mobilenetv3_small_minimal_100 |        ✅        |
+----------------------------------+------------------+
| tinynet_a                        |        ✅        |
+----------------------------------+------------------+
| tinynet_b                        |        ✅        |
+----------------------------------+------------------+
| tinynet_c                        |        ✅        |
+----------------------------------+------------------+
| tinynet_d                        |        ✅        |
+----------------------------------+------------------+
| tinynet_e                        |        ✅        |
+----------------------------------+------------------+
| vgg11                            |                  |
+----------------------------------+------------------+
| vgg11_bn                         |                  |
+----------------------------------+------------------+
| vgg13                            |                  |
+----------------------------------+------------------+
| vgg13_bn                         |                  |
+----------------------------------+------------------+
| vgg16                            |                  |
+----------------------------------+------------------+
| vgg16_bn                         |                  |
+----------------------------------+------------------+
| vgg19                            |                  |
+----------------------------------+------------------+
| vgg19_bn                         |                  |
+----------------------------------+------------------+
| vovnet39a                        |                  |
+----------------------------------+------------------+
| vovnet57a                        |                  |
+----------------------------------+------------------+
| wide_resnet50_2                  |        ✅        |
+----------------------------------+------------------+
| wide_resnet101_2                 |        ✅        |
+----------------------------------+------------------+
| xception41                       |        ✅        |
+----------------------------------+------------------+
| xception41p                      |        ✅        |
+----------------------------------+------------------+
| xception65                       |        ✅        |
+----------------------------------+------------------+
| xception65p                      |        ✅        |
+----------------------------------+------------------+
| xception71                       |        ✅        |
+----------------------------------+------------------+

Transformer-Style

All segmentation models in SMP (this library short name) are made of:

- encoder (feature extractor, a.k.a backbone)
- decoder (features fusion block to create segmentation *mask*)
- segmentation head (final head to reduce number of channels from decoder and upsample mask to preserve input-output spatial resolution identity)
- classification head (optional head which build on top of deepest encoder features)


2. Creating your own encoder

class MyEncoder(torch.nn.Module, EncoderMixin):
    
    def __init__(self, **kwargs):
        super().__init__()
        
        # A number of channels for each encoder feature tensor, list of integers
        self._out_channels: List[int] = [3, 16, 64, 128, 256, 512]

        # A number of stages in decoder (in other words number of downsampling operations), integer
        # use in in forward pass to reduce number of returning features
        self._depth: int = 5 

        # Default number of input channels in first Conv2d layer for encoder (usually 3)
        self._in_channels: int = 3 
        
        # Define encoder modules below
        ...

    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
        """Produce list of features of different spatial resolutions, each feature is a 4D torch.tensor of
        shape NCHW (features should be sorted in descending order according to spatial resolution, starting
        with resolution same as input `x` tensor).

        Input: `x` with shape (1, 3, 64, 64)
        Output: [f0, f1, f2, f3, f4, f5] - features with corresponding shapes
                [(1, 3, 64, 64), (1, 64, 32, 32), (1, 128, 16, 16), (1, 256, 8, 8),
                (1, 512, 4, 4), (1, 1024, 2, 2)] (C - dim may differ)

        also should support number of features according to specified depth, e.g. if depth = 5,
        number of feature tensors = 6 (one with same resolution as input and 5 downsampled),
        depth = 3 -> number of feature tensors = 4 (one with same resolution as input and 3 downsampled).
        """

        return [feat1, feat2, feat3, feat4, feat5, feat6]

smp.encoders.encoders["my_awesome_encoder"] = {
    "encoder": MyEncoder, # encoder class here
    "pretrained_settings": {
        "imagenet": {
            "mean": [0.485, 0.456, 0.406],
            "std": [0.229, 0.224, 0.225],
            "url": "https://some-url.com/my-model-weights",
            "input_space": "RGB",
            "input_range": [0, 1],
        },
    },
    "params": {
        # init params for encoder if any
    },
},

model = smp.Unet(encoder_name="my_awesome_encoder")

If it works fine, don`t forget to contribute your work and make a PR to SMP 😉

All models support ``aux_params`` parameter, which is default set to ``None``. 
If ``aux_params = None`` than classification auxiliary output is not created, else
model produce not only ``mask``, but also ``label`` output with shape ``(N, C)``.

Classification head consist of following layers:
    
1. GlobalPooling
2. Dropout (optional)
3. Linear
4. Activation (optional)

Example:

.. code-block:: python
    
    aux_params=dict(
        pooling='avg',             # one of 'avg', 'max'
        dropout=0.5,               # dropout ratio, default is None
        activation='sigmoid',      # activation function, default is None
        classes=4,                 # define number of output labels
    )

    model = smp.Unet('resnet34', classes=4, aux_params=aux_params)
    mask, label = model(x)

    mask.shape, label.shape
    # (N, 4, H, W), (N, 4)

4. Freezing and unfreezing the encoder

model = smp.Unet("resnet34", classes=2)

# Freeze encoder: stops gradient updates and freezes normalization layer stats
model.freeze_encoder()

# Unfreeze encoder: re-enables training for encoder parameters and normalization layers
model.unfreeze_encoder()

$ pip install -U segmentation-models-pytorch

$ pip install -U git+https://github.com/qubvel/segmentation_models.pytorch

.. automodule:: segmentation_models_pytorch.losses.constants
        :members:

JaccardLoss

.. autoclass:: segmentation_models_pytorch.losses.DiceLoss

TverskyLoss

.. autoclass:: segmentation_models_pytorch.losses.FocalLoss

LovaszLoss

.. autoclass:: segmentation_models_pytorch.losses.SoftBCEWithLogitsLoss

SoftCrossEntropyLoss
~~~~~~~~~~~~~~~~~~~~
.. autoclass:: segmentation_models_pytorch.losses.SoftCrossEntropyLoss

MCCLoss
~~~~~~~~~~~~~~~~~~~~
.. autoclass:: segmentation_models_pytorch.losses.MCCLoss
        :members: forward

---

### Metrics (`docs/metrics.rst`)

📏 Metrics
==========

Functional metrics
~~~~~~~~~~~~~~~~~~

.. currentmodule:: segmentation_models_pytorch.metrics.functional

.. autosummary::

   get_stats
   fbeta_score
   f1_score
   iou_score
   accuracy
   precision
   recall
   sensitivity
   specificity
   balanced_accuracy
   positive_predictive_value
   negative_predictive_value
   false_negative_rate
   false_positive_rate
   false_discovery_rate
   false_omission_rate
   positive_likelihood_ratio
   negative_likelihood_ratio

.. automodule:: segmentation_models_pytorch.metrics.functional
   :members:

---

### Models (`docs/models.rst`)

🕸️ Segmentation Models
==============================


.. contents::
   :local:

.. _unet:

Unet
~~~~
.. autoclass:: segmentation_models_pytorch.Unet


.. _unetplusplus:

Unet++
~~~~~~
.. autoclass:: segmentation_models_pytorch.UnetPlusPlus


.. _fpn:

FPN
~~~
.. autoclass:: segmentation_models_pytorch.FPN


.. _pspnet:

PSPNet
~~~~~~
.. autoclass:: segmentation_models_pytorch.PSPNet


.. _deeplabv3:

DeepLabV3
~~~~~~~~~
.. autoclass:: segmentation_models_pytorch.DeepLabV3


.. _deeplabv3plus:

DeepLabV3+
~~~~~~~~~~
.. autoclass:: segmentation_models_pytorch.DeepLabV3Plus


.. _linknet:

Linknet
~~~~~~~
.. autoclass:: segmentation_models_pytorch.Linknet


.. _manet:

MAnet
~~~~~~
.. autoclass:: segmentation_models_pytorch.MAnet


.. _pan:

PAN
~~~
.. autoclass:: segmentation_models_pytorch.PAN


.. _upernet:

UPerNet
~~~~~~~
.. autoclass:: segmentation_models_pytorch.UPerNet


.. _segformer:

Segformer
~~~~~~~~~
.. autoclass:: segmentation_models_pytorch.Segformer


.. _dpt:

DPT
~~~

.. note::

    See full list of DPT-compatible timm encoders in :ref:`dpt-encoders`.

.. note::

    For some encoders, the model requires ``dynamic_img_size=True`` to be passed in order to work with resolutions different from what the encoder was trained for.

.. autoclass:: segmentation_models_pytorch.DPT

---

### Quickstart (`docs/quickstart.rst`)

🚀 Quick Start
==============

**1. Create segmentation model**

Segmentation model is just a PyTorch nn.Module, which can be created as easy as:

.. code-block:: python

    import segmentation_models_pytorch as smp

    model = smp.Unet(
        encoder_name="resnet34",        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
        encoder_weights="imagenet",     # use `imagenet` pre-trained weights for encoder initialization
        in_channels=1,                  # model input channels (1 for gray-scale images, 3 for RGB, etc.)
        classes=3,                      # model output channels (number of classes in your dataset)
    )

- Check the page with available :doc:`model architectures <models>`.
- Check the table with :doc:`available ported encoders and its corresponding weights <encoders>`.
- `Pytorch Image Models (timm) <https://github.com/huggingface/pytorch-image-models>`_ encoders are also supported, check it :doc:`here<encoders_timm>`.

Alternatively, you can use `smp.create_model` function to create a model by name:

.. code-block:: python

    model = smp.create_model(
        arch="fpn",                     # name of the architecture, e.g. 'Unet'/ 'FPN' / etc. Case INsensitive!
        encoder_name="mit_b0",
        encoder_weights="imagenet",
        in_channels=1,
        classes=3,
    )


**2. Configure data preprocessing**

All encoders have pretrained weights. Preparing your data the same way as during weights pre-training may give your better results (higher metric score and faster convergence). But it is relevant only for 1-2-3-channels images and **not necessary** in case you train the whole model, not only decoder.

.. code-block:: python

    from segmentation_models_pytorch.encoders import get_preprocessing_fn

    preprocess_input = get_preprocessing_fn('resnet18', pretrained='imagenet')


**3. Congratulations!** 🎉


You are done! Now you can train your model with your favorite framework, or as simple as:

.. code-block:: python

    for images, gt_masks in dataloader:

        predicted_mask = model(images)
        loss = loss_fn(predicted_mask, gt_masks)

        loss.backward()
        optimizer.step()

Check the following examples:

.. |colab-badge| image:: https://colab.research.google.com/assets/colab-badge.svg
   :target: https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/main/examples/binary_segmentation_intro.ipynb
   :alt: Open In Colab

- Finetuning notebook on Oxford Pet dataset with `PyTorch Lightning <https://github.com/qubvel/segmentation_models.pytorch/blob/main/examples/binary_segmentation_intro.ipynb>`__ |colab-badge|
- Finetuning script for cloth segmentation with `PyTorch Lightning <https://github.com/ternaus/cloths_segmentation>`__

---

### Save Load (`docs/save_load.rst`)

📂 Saving and Loading
=====================

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.

Saving and Sharing a Model
--------------------------

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.
And, optionally, you can push the model to the Hugging Face Hub by setting the `push_to_hub` parameter to `True`.

For example:

.. code:: python

    import segmentation_models_pytorch as smp

    model = smp.Unet('resnet34', encoder_weights='imagenet')

    # After training your model, save it to a directory
    model.save_pretrained('./my_model')

    # Or saved and pushed to the Hub simultaneously
    model.save_pretrained('username/my-model', push_to_hub=True)

Loading Trained Model
---------------------

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.

For example:

.. code:: python

    import segmentation_models_pytorch as smp

    # Load the model from the local directory
    model = smp.from_pretrained('./my_model')

    # Alternatively, load the model directly from the Hugging Face Hub
    model = smp.from_pretrained('username/my-model')

Loading pre-trained model with different number of classes for fine-tuning:

.. code:: python

    import segmentation_models_pytorch as smp

    model = smp.from_pretrained('', classes=5, strict=False)

Saving model Metrics and Dataset Name
-------------------------------------

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.

For example:

.. code:: python

    import segmentation_models_pytorch as smp

    model = smp.Unet('resnet34', encoder_weights='imagenet')

    # After training your model, save it to a directory
    model.save_pretrained('./my_model', metrics={'accuracy': 0.95}, dataset='my_dataset')

    # Or saved and pushed to the Hub simultaneously
    model.save_pretrained('username/my-model', push_to_hub=True, metrics={'accuracy': 0.95}, dataset='my_dataset')

Saving with preprocessing transform (Albumentations)
----------------------------------------------------

You can save the preprocessing transform along with the model and push it to the Hub.
This can be useful when you want to share the model with the preprocessing transform that was used during training,
to make sure that the inference pipeline is consistent with the training pipeline.

.. code:: python

    import albumentations as A
    import segmentation_models_pytorch as smp

    # Define a preprocessing transform for image that would be used during inference
    preprocessing_transform = A.Compose([
        A.Resize(256, 256),
        A.Normalize()
    ])

    model = smp.Unet()

    directory_or_repo_on_the_hub = "qubvel-hf/unet-with-transform"  # /<repo-name>

    # Save the model and transform (and pus ot hub, if needed)
    model.save_pretrained(directory_or_repo_on_the_hub, push_to_hub=True)
    preprocessing_transform.save_pretrained(directory_or_repo_on_the_hub, push_to_hub=True)

    # Loading transform and model
    restored_model = smp.from_pretrained(directory_or_repo_on_the_hub)
    restored_transform = A.Compose.from_pretrained(directory_or_repo_on_the_hub)

    print(restored_transform)

Conclusion
----------

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.

|binary-segmentation-intro|
|save-load-model-and-share-with-hf-hub|

.. |binary-segmentation-intro| image:: https://colab.research.google.com/assets/colab-badge.svg
    :target: https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/main/examples/binary_segmentation_intro.ipynb
    :alt: Open In Colab

.. |save-load-model-and-share-with-hf-hub| image:: https://colab.research.google.com/assets/colab-badge.svg
    :target: https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/main/examples/save_load_model_and_share_with_hf_hub.ipynb
    :alt: Open In Colab

---

--- METRICS ---
- Files Extracted: 12
- Estimated Token Budget: ~38668 tokens
- Canonical Reference: https://codewiki.google/github.com/qubvel-org/segmentation_models.pytorch