https://github.com/lucidrains/segformer-pytorch
Implementation of Segformer, Attention + MLP neural network for segmentation, in Pytorch
https://github.com/lucidrains/segformer-pytorch
artificial-intelligence attention-mechanism deep-learning image-segmentation multilayer-perceptron segmentation
Last synced: 6 months ago
JSON representation
Implementation of Segformer, Attention + MLP neural network for segmentation, in Pytorch
- Host: GitHub
- URL: https://github.com/lucidrains/segformer-pytorch
- Owner: lucidrains
- License: mit
- Created: 2021-06-06T19:06:49.000Z (over 4 years ago)
- Default Branch: main
- Last Pushed: 2022-12-09T02:37:48.000Z (almost 3 years ago)
- Last Synced: 2025-03-29T03:03:21.168Z (6 months ago)
- Topics: artificial-intelligence, attention-mechanism, deep-learning, image-segmentation, multilayer-perceptron, segmentation
- Language: Python
- Homepage:
- Size: 104 KB
- Stars: 358
- Watchers: 9
- Forks: 45
- Open Issues: 11
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
## Segformer - Pytorch
Implementation of Segformer, Attention + MLP neural network for segmentation, in Pytorch.
## Install
```bash
$ pip install segformer-pytorch
```
## Usage
For example, MiT-B0
```python
import torch
from segformer_pytorch import Segformer
model = Segformer(
dims = (32, 64, 160, 256), # dimensions of each stage
heads = (1, 2, 5, 8), # heads of each stage
ff_expansion = (8, 8, 4, 4), # feedforward expansion factor of each stage
reduction_ratio = (8, 4, 2, 1), # reduction ratio of each stage for efficient attention
num_layers = 2, # num layers of each stage
decoder_dim = 256, # decoder dimension
num_classes = 4 # number of segmentation classes
)
x = torch.randn(1, 3, 256, 256)
pred = model(x) # (1, 4, 64, 64) # output is (H/4, W/4) map of the number of segmentation classes
```
Make sure the keywords are at most a tuple of 4, as this repository is hard-coded to give the MiT 4 stages as done in the paper.
## Citations
```bibtex
@misc{xie2021segformer,
title = {SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers},
author = {Enze Xie and Wenhai Wang and Zhiding Yu and Anima Anandkumar and Jose M. Alvarez and Ping Luo},
year = {2021},
eprint = {2105.15203},
archivePrefix = {arXiv},
primaryClass = {cs.CV}
}
```