登录

segmentation_models.pytorch

qubvel-orgMIT活跃维护

基于 PyTorch 的语义分割模型库，提供 500+ 预训练卷积与 Transformer 骨干网络，开箱即用，省去自己搭建和训练分割模型的麻烦。

Semantic segmentation models with 500+ pretrained convolutional and transformer-based backbones.

v0.6.0发布于 2026-10-26 16:43 UTC11,801star +3 / 7 天1,853fork同步于 2026-11-10 11:28 UTC
语义分割预训练模型PyTorch骨干网络图像分割
对比 GitHub 官网computer-visiondeeplab-v3-plusdeeplabv3dptfpn

怎么上手

由 README 提炼，执行前请确认
原文 
pip install segmentation-models-pytorch
硬件要求按 README 说明

README 未说明具体显存或 GPU 型号要求，仅标注依赖 PyTorch 1.9+ 与 Python 3.9+

下载：源码 zip·GitHub Releases（安装包）本站不镜像安装包

项目说明

由 README 自动翻译提炼
没人工点评？写下你的经验

基于 PyTorch 的语义分割模型库，两行代码即可搭建带预训练骨干的分割网络

基于 PyTorch 的语义分割模型库，提供 500+ 预训练卷积与 Transformer 骨干网络，开箱即用，省去自己搭建和训练分割模型的麻烦。

核心能力

  • 两行高层 API 创建 Unet 等分割模型
  • 内置 12 种编码器-解码器架构如 Unet++、Segformer、DPT
  • 提供 800+ 预训练卷积与 Transformer 骨干并支持 timm
  • 附带 Dice、Jaccard、Tversky 等常用损失与指标
  • 支持 ONNX 导出及 torch script/trace/compile

容易踩的坑

首次使用预训练权重需联网下载对应骨干权重；需自行安装匹配版本的 PyTorch（1.9+）与 Python（3.9+）；使用 timm 骨干需额外安装 timm

内容由模型根据 README 生成，可能不完整或不准确，以官方文档为准。看原文 README

最近发版

5 个版本v0.6.0
全部 release
  1. v0.6.0Segmentation Models - v0.6.0最新正式版2026-10-26 · 2026-10-26 16:43 UTC

    ⚠️ Breaking changes

    • Python 3.10+ is required, and the minimum dependency versions went up: torch>=1.11, torchvision>=0.12, numpy>=1.21.2, pillow>=8.4 (#1281).
    • Encoder weights are downloaded only from the Hugging Face Hub. The old fallback to the original download URLs is gone, along with smp.encoders.pretrained_settings. If you can't reach the Hub, set HF_ENDPOINT to a mirror or load weights from a local file (#1398).
    • tu- encoders now warn when encoder_weights is a string. Use True or None, and choose the pretrained variant in the name, e.g. tu-convnext_tiny.fb_in22k (#1269).

    Highlights

    Freeze and unfreeze the encoder

    model = smp.Unet("resnet34")
model.freeze_encoder()    # also stops BatchNorm running-stat updates
...
model.unfreeze_encoder()
    • Add encoder_freeze method to SegmentationModel by @ricber in #1220

    Losses

    • class_weights for DiceLoss, JaccardLoss, TverskyLoss and FocalLoss (multiclass and multilabel modes): 
      loss = smp.losses.DiceLoss("multiclass", class_weights=[0.2, 1.0, 3.0])
    • FocalLoss accepts probabilities via from_logits=False, with an eps parameter like DiceLoss.
    • JaccardLoss supports ignore_index.
    • Losses work on non-contiguous tensors.

    Models and encoders

    • Transformer-style tu- encoders work in UnetPlusPlus and Linknet (#1294).
    • PSPNet exports to ONNX (#1270).
    • Fixed an MPS float tensor ValueError (#1199).

    What's Changed

  2. v0.5.0Segmentation Models - v0.5.02025-04-17 · 2025-04-17 10:13 UTC
    查看更新说明

    New Models

    DPT

    DPT

    The DPT model adapts the Vision Transformer (ViT) architecture for dense prediction tasks like semantic segmentation. It uses a ViT as a powerful backbone, processing image information with a global receptive field at each stage. The key innovation lies in its decoder, which reassembles token representations from various transformer stages into image-like feature maps at different resolutions. These are progressively combined using convolutional PSP and FPN blocks to produce full-resolution, high-detail predictions.

    The model in smp can be used with a wide variety of transformer-based encoders

    import segmentation_models_pytorch as smp

# initialize with your own pretrained encoder
model = smp.DPT("tu-mobilevitv2_175.cvnets_in1k", classes=2)

# load fully-pretrained on ADE20K 
model = smp.from_pretrained("smp-hub/dpt-large-ade20k")

# load the same checkpoint for finetuning
model = smp.from_pretrained("smp-hub/dpt-large-ade20k", classes=1, strict=False)

    The full table of DPT's supported timm encoders can be found here.

    Models export

    A lot of work was done to add support for torch.jit.script, torch.compile (without graph breaks: fullgraph=True) and torch.export features in all encoders and models.

    This provides several advantages:

    • torch.jit.script: Enables serialization of models into a static graph format, enabling deployment in environments without a Python interpreter and allowing for graph-based optimizations.
    • torch.compile (with fullgraph=True): Leverages Just-In-Time (JIT) compilation (e.g., via Triton or Inductor backends) to generate optimized kernels, reducing Python overhead and enabling significant performance improvements through techniques like operator fusion, especially on GPU hardware. fullgraph=True minimizes graph breaks, maximizing the scope of these optimizations.
    • torch.export: Produces a standardized Ahead-Of-Time (AOT) graph representation, simplifying the process of exporting models to various inference backends and edge devices (e.g., through ExecuTorch) while preserving model dynamism where possible.

    PRs:

    Core

    All encoders from third-party libraries such as efficientnet-pytorch and pretrainedmodels.pytorch are now vendored by SMP. This means we have copied and refactored the underlying code and moved all checkpoints to the smp-hub. As a result, you will have fewer additional dependencies when installing smp and get much faster weights downloads.

    🚨🚨🚨 Breaking changes

    1. UperNet model was significantly changed to reflect the original implementation and to bring pretrained checkpoints into SMP. Unfortunately, UperNet model weights trained with v0.4.0 will be not compatible with SMP v0.5.0.

    2. While the high-level API for modeling should be backward compatible with v0.4.0, internal modules (such as encoders, decoders

  3. v0.4.0Segmentation Models - v0.4.02025-01-08 · 2025-01-08 15:28 UTC
    查看更新说明

    New models

    Segformer

    contributed by @brianhou0208

    Originally, SegFormer is a transformer-based semantic segmentation model known for its simplicity and efficiency. It uses a lightweight hierarchical encoder to capture multi-scale features and a minimal decoder for fast inference.

    With segmentation-models-pytorch you can utilize the model with a native Mix Vision Transformer encoder as long as with 800+ other encoders supported by the library. Original weights are also supported and can be loaded as follows:

    import segmentation_models_pytorch as smp

model = smp.from_pretrained("smp-hub/segformer-b5-640x640-ade-160k")

    or with any other encoder:

    import segmentation_models_pytorch as smp

model = smp.Segformer("resnet34")

    See more checkpoints on the HF Hub.

    UperNet

    contributed by @brianhou0208

    UPerNet (Unified Perceptual Parsing Network) is a versatile semantic segmentation model designed to handle diverse scene parsing tasks. It combines a Feature Pyramid Network (FPN) with a Pyramid Pooling Module (PPM) to effectively capture multi-scale context.

    import segmentation_models_pytorch as smp

model = smp.UPerNet("resnet34")

    New Encoders

    Thanks to @brianhou0208 contribution 800+ timm encoders are now supported in segmentation_models.pytorch. New modern encoders like convnext, efficientvit, efficientformerv2, hiera, mambaout and more can be used as easy as:

    import segmentation_models_pytorch as smp

model = smp.create_model("upernet", encoder_name="tu-mambaout_small")
# or
model = smp.UPerNet("tu-mambaout_small")

    New examples

    • Added example for multi-class segmentation by @TimbusCalin
    • Added example for onnx export by @qubvel

    Other changes

    • Project migrated to pyproject.toml by @adamjstewart
    • Better dependency managing and testing (minimal and latest dependencies, linux/windows/mac platforms) by @adamjstewart
    • Better type annotations
    • Tests are refactored for faster CI and local testing by @qubvel

    All changes

  4. v0.3.4Segmentation Models - v0.3.42024-08-23 · 2024-08-23 13:15 UTC
    查看更新说明

    Updates

    • 🤗 Hugging Face integration: you can save, load, and share models with HF Hub, see example notebook.

    Full log

    New Contributors

    Full Changelog: https://github.com/qubvel-org/segmentation_models.pytorch/compare/v0.3.3...v0.3.4

  5. v0.3.3Segmentation Models - v0.3.32023-05-28 · 2023-05-28 15:49 UTC
    查看更新说明

    Updates

    • Pytorch image models (timm) version upgrade to 0.9.2

用户点评

登录用户的真实使用记录，可按版本和硬件筛选
共 0 条

    还没有人点评。第一个说出真实体验的人最有价值。

    README 原文

    未翻译
    在 GitHub 打开

    logo Python library with Neural Networks for Image Semantic Segmentation based on PyTorch.

    GitHub Workflow Status (branch) Codecov Read the Docs PyPI PyTorch - Version Python - Version Generic badge PyPI - Downloads

    The main features of the library are:

    • Super simple high-level API (just two lines to create a neural network)
    • 12 encoder-decoder model architectures (Unet, Unet++, Segformer, DPT, ...)
    • 800+ pretrained convolution- and transform-based encoders, including timm support
    • Popular metrics and losses for training routines (Dice, Jaccard, Tversky, ...)
    • ONNX export and torch script/trace/compile friendly

    🤝 Sponsor: withoutBG

    withoutBG is a high-quality background removal tool. They built their open-source image matting and refiner models using smp.Unet and are proudly sponsoring this project.

    📚 Project Documentation 📚

    Visit Read The Docs Project Page or read the following README to know more about Segmentation Models Pytorch (SMP for short) library

    📋 Table of content

    1. Quick start
    2. Examples
    3. Models and encoders
    4. Models API
      1. Input channels
      2. Auxiliary classification output
      3. Depth
    5. Installation
    6. Competitions won with the library
    7. Contributing
    8. Citing
    9. License

    ⏳ Quick start

    1. Create your first Segmentation model with SMP

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

    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)
)
    • see table with available model architectures
    • see table with available encoders and their corresponding weights
    2. Configure data preprocessing

    All encoders have pretrained weights. Preparing your data the same way as during weights pre-training may give you better results (higher metric score and faster convergence). It is not necessary in case you train the whole model, not only the decoder.

    from segmentation_models_pytorch.encoders import get_preprocessing_fn

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

    Congratulations! You are done! Now you can train your model with your favorite framework!

    💡 Examples

    NameLinkColab
    Train pets binary segmentation on OxfordPetsNotebookOpen In Colab
    Train cars binary segmentation on CamVidNotebookOpen In Colab
    Train multiclass segmentation on CamVidNotebookOpen In Colab
    Train clothes binary segmentation by @ternausRepo
    Load and inference pretrained SegformerNotebookOpen In Colab
    Load and inference pretrained DPTNotebookOpen In Colab
    Load and inference pretrained UPerNetNotebookOpen In Colab
    Save and load models locally / to HuggingFace HubNotebookOpen In Colab
    Export trained model to ONNXNotebookOpen In Colab

    📦 Models and encoders

    Architectures

    ArchitecturePaperDocumentationCheckpoints
    Unetpaperdocs
    Unet++paperdocs
    MAnetpaperdocs
    Linknetpaperdocs
    FPNpaperdocs
    PSPNetpaperdocs
    PANpaperdocs
    DeepLabV3paperdocs
    DeepLabV3+paperdocs
    UPerNetpaperdocscheckpoints
    Segformerpaperdocscheckpoints
    DPTpaperdocscheckpoints

    Encoders

    The library provides a wide range of pretrained encoders (also known as backbones) for segmentation models. Instead of using features from the final layer of a classification model, we extract intermediate features and feed them into the decoder for segmentation tasks.

    All encoders come with pretrained weights, which help achieve faster and more stable convergence when training segmentation models.

    Given the extensive selection of supported encoders, you can choose the best one for your specific use case, for example:

    • Lightweight encoders for low-latency applications or real-time inference on edge devices (mobilenet/mobileone).
    • High-capacity architectures for complex tasks involving a large number of segmented classes, providing superior accuracy (convnext/swin/mit).

    By selecting the right encoder, you can balance efficiency, performance, and model complexity to suit your project needs.

    All encoders and corresponding pretrained weight are listed in the documentation:

    • table with natively ported encoders
    • table with timm encoders supported

    🔁 Models API

    Input channels

    The input channels parameter allows you to create a model that can process a tensor with an arbitrary number of channels. If you use pretrained weights from ImageNet, the weights of the first convolution will be reused:

    • For the 1-channel case, it would be a sum of the weights of the first convolution layer.
    • Otherwise, channels would be populated with weights like new_weight[:, i] = pretrained_weight[:, i % 3], and then scaled with new_weight * 3 / new_in_channels.
    model = smp.FPN('resnet34', in_channels=1)
mask = model(torch.ones([1, 1, 64, 64]))

    Auxiliary classification output

    All models support aux_params parameters, which is default set to None. If aux_params = None then classification auxiliary output is not created, else model produce not only mask, but also label output with shape NC. Classification head consists of GlobalPooling->Dropout(optional)->Linear->Activation(optional) layers, which can be configured by aux_params as follows:

    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)

    Depth

    Depth parameter specify a number of downsampling operations in encoder, so you can make your model lighter if specify smaller depth.

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

    🛠 Installation

    PyPI version:

    $ pip install segmentation-models-pytorch

    The latest version from GitHub:

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

    🏆 Competitions won with the library

    Segmentation Models package is widely used in image segmentation competitions. Here you can find competitions, names of the winners and links to their solutions.

    🛠 Projects built with SMP

    🤝 Contributing

    1. Install SMP in dev mode
    make install_dev  # Create .venv, install SMP in dev mode
    1. Run tests and code checks
    make test          # Run tests suite with pytest
make fixup         # Ruff for formatting and lint checks
    1. Update a table (in case you added an encoder)
    make table        # Generates a table with encoders and print to stdout

    📝 Citing 

    @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}}
}

    🛡️ License

    The project is primarily distributed under MIT License, while some files are subject to other licenses. Please refer to LICENSES and license statements in each file for careful check, especially for commercial use.

    相关项目

    同分类里 star 最高的项目
    对比全部 →
    opencv
    opencv·C++

    开源计算机视觉库，提供图像处理、视频分析和深度学习推理的 C++ 接口，是视觉应用开发的基础设施。

    计算机视觉+4图像处理
    91517+4157k5.1.02026-11-10 10:53 UTC
    upscayl
    upscayl·TypeScript

    跨平台的开源 AI 图片放大工具，基于 ESRGAN 等模型把低分辨率图片无损放大并增强细节，有图形界面，普通用户也能直接上手。

    图像超分+4图像处理
    52145+2652.7kv2.15.02026-10-05 19:13 UTC
    supervision
    roboflow·Python

    一套可复用的计算机视觉工具库，提供目标检测、实例分割、分类等任务的标注处理、指标计算与可视化组件，帮开发者省去重复造轮子的时间。

    计算机视觉+4图像处理
    51577+444.9k0.31.52026-11-10 11:15 UTC
    sharp
    lovell·JavaScript

    基于 libvips 的高性能 Node.js 图像处理库，能快速完成 JPEG、PNG、WebP、AVIF、TIFF 的缩放裁剪与格式转换，适合服务端批量处理图片。

    图像处理+4图像处理
    32836+71.4kv0.35.52026-11-09 13:05 UTC