Package: torchseg
TorchSegTorchSeg is an actively maintained and up-to-date fork of theSegmentation Models PyTorch (smp) library.UpdatesThe goal of this fork is to 1) provide maintenance support for the original library and 2) add features relevant to modern semantic segmentation. Since the fork, this library has added some features which can be summarized below:ImprovedPyTorch Image Models (timm)for models with feature extraction functionality (852/1017=84% of timm models). This includes the typical CNN models such asResNet,EfficientNet, etc., but now extends to include modern architectures likeConvNext,Swin,PoolFormer,MaxViTand more!Support for pretrained Vision Transformer (ViT) encoders. Currently timm ViTs do not support feature extraction out of the box. However we have added support for extracting intermediate transformer encoder layer feature maps to obtain this functionality. We support 100+ ViT based models includingViT,DeiT,FlexiViT!Additionally we have performed the following for improved software standards:More thorough testing and CIFormatting usingblack,isort,flake8,mypyReduction of dependence on unmaintained libraries (now depends only ontorch,timm, andeinops)Reduce lines of code to maintain (removed custom utils, metrics, encoders) in favor of newer libraries such astorchmetricsandtimmFeaturesThe main features of this library are:High level API (just two lines to create a neural network)9 segmentation architectures for binary and multi class segmentation (including U-Net, DeepLabV3)Support for 852/1017 (~84%) of available encoders fromtimmAll encoders have pre-trained weights for faster and better convergencePopular segmentation loss functionsExample UsageTorchSeg models at their base are just torch nn.Modules. They can be created as follows:importtorchsegmodel=torchseg.Unet(encoder_name="resnet50",encoder_weights=True,in_channels=3classes=3,)TorchSeg has anencoder_paramsfeature which passes additional parameters totimm.create_model()when defining an encoder backbone. One can specify different activitions, normalization layers, and more like below.You can also define afunctools.partialcallable as an activation/normalization layer. See the timm docs for more information on availableactivationsandnormalizationlayers. You can even used pretrained weights while changing the activations/normalizations!model=torchseg.Unet(encoder_name="resnet50",encoder_weights=True,in_channels=3classes=3,encoder_params={"act_layer":"prelu","norm_layer":"layernorm"})Some models likeSwinandConvNextperform a downsampling of scale=4 in the first block (stem) and then downsample by 2 afterwards with onlydepth=4blocks. This results in an output size of half after the decoder. To get the same output size as the input you can passhead_upsampling=2which will upsample once more prior to the segmentation head.model=torchseg.Unet("convnextv2_tiny",in_channels=3,classes=2,encoder_weights=True,encoder_depth=4,decoder_channels=(256,128,64,32),head_upsampling=2)model=torchseg.Unet("swin_tiny_patch4_window7_224",in_channels=3,classes=2,encoder_weights=True,encoder_depth=4,decoder_channels=(256,128,64,32),head_upsampling=2,encoder_params={"img_size":256}# need to define img size since swin is a ViT hybrid)model=torchseg.Unet("maxvit_small_tf_224",in_channels=3,classes=2,encoder_weights=True,encoder_depth=5,decoder_channels=(256,128,64,32,16),encoder_params={"img_size":256})TorchSeg supports pretrained ViT encoders from timm by extracting intermediate transformer block features specified by theencoder_indicesandencoder_deptharguments.You will also need to definescale_factorsfor upsampling the feature layers to the resolutions expected by the decoders. For U-Netdepth=5this would bescales=(8, 4, 2, 1, 0.5). Fordepth=4this would bescales=(4, 2, 1, 0.5), fordepth=3this would bescales=(2, 1, 0.5)and so on.Another benefit of using timm is that by passing in a newimg_size, timm automatically interpolates the ViT positional embeddings to work with your new image size which creates a different number of patch tokens.importtorchimporttorchsegmodel=torchseg.Unet("vit_small_patch16_224",in_channels=8,classes=2,encoder_depth=5,encoder_indices=(2,4,6,8,10),# which intermediate blocks to extract features fromencoder_weights=True,decoder_channels=(256,128,64,32,16),encoder_params={# additional params passed to timm.create_model and the vit encoder"scale_factors":(8,4,2,1,0.5),# resize scale_factors for patch size 16 and 5 layers"img_size":256,# timm automatically interpolates the positional embeddings to your new image size},)ModelsArchitectures (Decoders)Unet [paper]Unet++ [paper]MAnet [paper]Linknet [paper]FPN [paper]PSPNet [paper]PAN [paper]DeepLabV3 [paper]DeepLabV3+ [paper]EncodersTorchSeg relies entirely on thetimmlibrary for pretrained encoder support. This means that TorchSeg supports any timm model which hasfeatures_onlyfeature extraction functionality. Additionally we support any ViT models with aget_intermediate_layersmethod. This results in a total of 852/1017 (~84%) encoders from timm includingResNet,Swin,ConvNext,ViT, and more!To list the following supported encoders:importtorchsegtorchseg.list_encoders()We have additionally pulled the the feature extractor metadata of each model withfeatures_onlysupport from timm atoutput_stride=32. This metadata provides information such as the number of intermediate layers, channels for each layer, layer name, and downsampling reduction.importtorchsegmetadata=torchseg.encoders.TIMM_ENCODERS["convnext_base"]print(metadata)"""{'channels': [128, 256, 512, 1024],'indices': (0, 1, 2, 3),'module': ['stages.0', 'stages.1', 'stages.2', 'stages.3'],'reduction': [4, 8, 16, 32],}"""metadata=torchseg.encoders.TIMM_ENCODERS["resnet50"]print(metadata)"""{'channels': [64, 256, 512, 1024, 2048],'indices': (0, 1, 2, 3, 4),'module': ['act1', 'layer1', 'layer2', 'layer3', 'layer4'],'reduction': [2, 4, 8, 16, 32]}"""Models APImodel.encoder- pretrained backbone to extract intermediate featuresmodel.decoder- network for processing the intermediate features to the original image resolution (Unet,DeepLabv3+,FPN)model.segmentation_head- final block producing the mask output (includes optional upsampling and activation)model.classification_head- optional block which create classification head on top of encodermodel.forward(x)- sequentially passxthrough model`s encoder, decoder and segmentation head (and classification head if specified)Input channelsTimm encoders supports the use of pretrained weights with arbitrary input channels by repeating weights for channels if > 3. For example, ifin_channels=6, RGB ImageNet pretrained weights in the initial layer would be repeatedRGBRGBto avoid random initialization. Forin_channels=7this would result inRGBRGBR. Below is a diagram to visualize this method.Auxiliary ClassifierAll models support an optional auxiliary classifier head through the use ofaux_params. Ifaux_params != Nonethen the
model will produce the alabeloutput in addition to themaskoutput with shape(N, C).
Classification head consists of GlobalPooling->Dropout(optional)->Linear->Activation(optional) layers, which can be
configured byaux_paramsas follows:aux_params=dict(pooling='avg',# one of 'avg', 'max'dropout=0.5,# dropout ratio, default is Noneactivation=nn.Sigmoid(),# activation function, default is Identityclasses=4,# define number of output labels)model=torchseg.Unet('resnet18',classes=4,aux_params=aux_params)mask,label=model(x)DepthDepth represents the number of downsampling operations in the encoder, so you can make
your model lighter by specifying lessdepth. Defaults todepth=5.Note that some models likeConvNextandSwinonly have 4 intermediate feature blocks. Therefore, to use these encoders setencoder_depth=4. This can be found in the metadata above.model=torchseg.Unet('resnet50',encoder_depth=4)