DOCSEARCHbyAlgolia
Submit Docs

SegFormer

repository·master·Indexed 25 days ago

https://github.com/nvlabs/segformer

A simple, efficient, and powerful semantic segmentation method using Transformers. This official PyTorch implementation is built upon the MMSegmentation codebase and provides tools for training and evaluation, along with a wide array of pre-trained models for datasets including Cityscapes, ADE20K, and Pascal VOC 2012.

Tokens
28.9K
Snippets
59
Records
143
Agent score
83%

What's inside SegFormer

  2. Prepare a model for publishing

    master

    Use tools/publish_model.py to prepare model weights for upload (e.g., to AWS). This script performs the following:

    1. Converts model weights to CPU tensors.
    2. Deletes optimizer states.
    3. Computes the hash of the checkpoint file and appends the hash ID to the filename.

    Example usage:

    python tools/publish_model.py work_dirs/pspnet/latest.pth psp_r50_hszhao_200ep.pth

    If the input is psp_r50_hszhao_200ep.pth, the final output filename will include a hash ID, such as psp_r50_512x1024_40ki_cityscapes-{hash id}.pth.

    python tools/publish_model.py ${INPUT_FILENAME} ${OUTPUT_FILENAME}

  3. Customize the optimizer constructor

    master

    For fine-grained parameter tuning (e.g., different weight decay for specific layers like BatchNorm), you can implement a custom optimizer constructor. Use the @OPTIMIZER_BUILDERS.register_module decorator. The class should implement a __call__(self, model) method that returns the configured optimizer instance.

    from mmcv.utils import build_from_cfg
from mmcv.runner import OPTIMIZER_BUILDERS

@OPTIMIZER_BUILDERS.register_module
class CocktailOptimizerConstructor(object):
    def __init__(self, optimizer_cfg, paramwise_cfg=None):
        self.optimizer_cfg = optimizer_cfg
        self.paramwise_cfg = paramwise_cfg

    def __call__(self, model):
        # return configured optimizer
        return my_optimizer

  4. Generate PNG files for Cityscapes official evaluation

    master

    To generate PNG files for submission to the official Cityscapes evaluation server using multi-GPU testing, follow these steps:

    1. Update your config file (e.g., configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py) to include the test directory paths:
    data = dict(
    test=dict(
        img_dir='leftImg8bit/test',
        ann_dir='gtFine/test'))
    1. Run the distributed test script with the --format-only flag and specify an output prefix via --eval-options:
    ./tools/dist_test.sh configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py \
    checkpoints/pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-2966598c.pth \
    4 --format-only --eval-options "imgfile_prefix=./pspnet_test_results"
    1. The PNG files will be located in the ./pspnet_test_results directory. You can then zip this directory for submission.
    ./tools/dist_test.sh configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py \
    checkpoints/pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-2966598c.pth \
    4 --format-only --eval-options "imgfile_prefix=./pspnet_test_results"

  5. Add a new loss function

    master

    To add a custom loss function:

    1. Implement the loss logic in a new file (e.g., mmseg/models/losses/my_loss.py).
    2. Use the @weighted_loss decorator on the functional implementation to enable element-wise weighting.
    3. Define a class inheriting from torch.nn.Module decorated with @LOSSES.register_module.
    4. Implement __init__ and forward (accepting pred, target, weight, avg_factor, and reduction_override).
    5. Import the module in mmseg/models/losses/__init__.py.
    6. Use it in your config by setting the loss_decode field in the decode_head to dict(type='MyLoss', ...).
    import torch
import torch.nn as nn
from ..builder import LOSSES
from .utils import weighted_loss

@weighted_loss
def my_loss(pred, target):
    assert pred.size() == target.size() and target.numel() > 0
    loss = torch.abs(pred - target)
    return loss

@LOSSES.register_module
class MyLoss(nn.Module):
    def __init__(self, reduction='mean', loss_weight=1.0):
        super(MyLoss, self).__init__()
        self.reduction = reduction
        self.loss_weight = loss_weight

    def forward(self, pred, target, weight=None, avg_factor=None, reduction_override=None):
        reduction = reduction_override if reduction_override else self.reduction
        loss = self.loss_weight * my_loss(
            pred, target, weight, reduction=reduction, avg_factor=avg_factor)
        return loss

    In your config file:

    loss_decode=dict(type='MyLoss', loss_weight=1.0)

  6. Install SegFormer/MMSegmentation on Windows (Experimental)

    master

    Windows support is experimental. Ensure you have a native C++ compiler (like cl.exe) and that it is added to your %PATH%.

    1. Create and activate a conda environment:
    conda create -n open-mmlab python=3.7 -y
conda activate open-mmlab
    1. Install PyTorch and torchvision:
    conda install pytorch=1.6.0 torchvision cudatoolkit=10.1 -c pytorch
    1. Install MMCV:

    mmcv-full is not supported on Windows. Install mmcv via pip:

    pip install mmcv
    1. Install MMSegmentation:
    git clone https://github.com/open-mmlab/mmsegmentation.git
cd mmsegmentation
pip install -e .

    Important Windows Notes:

    • Replace all backslashes \ in paths with forward slashes / in your Python code (e.g., using .replace('\', '/')).
    • Use absolute paths when linking datasets via mklink.
    conda create -n open-mmlab python=3.7 -y
conda activate open-mmlab

conda install pytorch=1.6.0 torchvision cudatoolkit=10.1 -c pytorch
set PATH=full\path\to\your\cpp\compiler;%PATH%
pip install mmcv

git clone https://github.com/open-mmlab/mmsegmentation.git
cd mmsegmentation
pip install -e .

mklink /D data %DATA_ROOT%

  7. Add a new segmentation head

    master

    To implement a new segmentation decode head:

    1. Create a new file in mmseg/models/decode_heads/.
    2. Define a class that inherits from BaseDecodeHead and decorate it with @HEADS.register_module().
    3. Implement __init__, init_weights, and forward.
    4. Import the module in mmseg/models/decode_heads/__init__.py.
    5. Configure it in your model config under the decode_head field.
    from .decode_head import BaseDecodeHead
from ..registry import HEADS

@HEADS.register_module()
class PSPHead(BaseDecodeHead):
    def __init__(self, pool_scales=(1, 2, 3, 6), **kwargs):
        super(PSPHead, self).__init__(**kwargs)

    def init_weights(self):
        pass

    def forward(self, inputs):
        pass

    In your config file:

    model = dict( decode_head=dict( type='PSPHead', in_channels=2048, in_index=3, channels=512, pool_scales=(1, 2, 3, 6), num_classes=19, ...) )

  8. Prepare Pascal Context dataset

    master

    To use the Pascal Context dataset, first install the Detail API. Then, use the conversion script with the trainval_merged.json file to format the annotations correctly.

    python tools/convert_datasets/pascal_context.py data/VOCdevkit data/VOCdevkit/VOC2010/trainval_merged.json

  9. Apply different learning rates for backbone and heads

    master

    To improve performance or convergence, you can set a different learning rate for the model heads compared to the backbone. In MMSegmentation, use paramwise_cfg within the optimizer dictionary to apply a multiplier (lr_mult) to specific parameter groups based on their names. For example, to make the learning rate of the heads 10 times larger than the backbone, add custom_keys with 'head' to the configuration.

    optimizer=dict(
    paramwise_cfg = dict(
        custom_keys={
            'head': dict(lr_mult=10.)}))

  10. Run speed benchmarks

    master
    To compute the average inference time (including network forwarding and post-processing, but excluding data loading), use the tools/benchmark.py script. The benchmark is typically performed on 200 images with torch.backends.cudnn.benchmark=False to ensure consistency.

  11. Convert a PyTorch model to ONNX format

    master

    Use the experimental tools/pytorch2onnx.py script to convert a model to ONNX format. The resulting model can be visualized using tools like Netron. The script also supports verifying the output results by comparing the PyTorch and ONNX models.

    Note: This tool is experimental and some customized operators are currently unsupported.

    python tools/pytorch2onnx.py ${CONFIG_FILE} --checkpoint ${CHECKPOINT_FILE} --output-file ${ONNX_FILE} [--shape ${INPUT_SHAPE} --verify]

  12. Prepare Cityscapes dataset

    master

    After downloading the Cityscapes dataset, use the conversion script to generate **labelTrainIds.png files required for training. You can specify the number of processes using --nproc to speed up conversion.

    # --nproc means 8 process for conversion, which could be omitted as well.
python tools/convert_datasets/cityscapes.py data/cityscapes --nproc 8