VisionMamba (kyegomez): A PyTorch Vim Class With No Training Script
Implementation of Vision Mamba from the paper: "Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model" It's 2.8x faster than DeiT and saves 86.8% GPU memory when performing batch inference to extract features on high-res images
At a glance
- What is it?
- VisionMamba is a small MIT-licensed PyTorch package that wraps the Vim architecture from the Vision Mamba paper in a single importable class. The forward pass is documented; the training path is not, and the README still lists the ImageNet training script as an open item.
- Who is it for?
- Adopt this package if you want a pip-installable PyTorch module that instantiates the Vim architecture with a documented set of constructor arguments, and you are prepared to write your own training loop. Do not adopt it if you need a reproducible ImageNet training pipeline, pretrained weights, or published accuracy numbers: the README's own Todo list still has the ImageNet training script unchecked, and the repository ships no releases.
- Can I use it commercially?
- Yes. MIT is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
- Is it still maintained?
- Yes. The repository last received commits 17 days ago.
- What is it written in?
- Mainly Python, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
The gap VisionMamba fills: a pip-installable Vim module
The Vision Mamba paper describes a visual backbone that replaces self-attention with a bidirectional state space model. Reading a paper gives you the equations. It does not give you a Python class you can put inside a training loop. VisionMamba exists to close that gap: the README states it is an implementation of Vision Mamba from the paper, and the installation instruction is a single pip command. The intended user is a PyTorch engineer who has read the paper, wants the architecture available as an importable module, and is willing to supply the surrounding training code. It is not aimed at someone who wants a finished image classifier with downloadable checkpoints. The README's headline claim, that the architecture is 2.8x faster than DeiT and saves 86.8% GPU memory when performing batch inference to extract features on high-res images, is quoted from the paper's framing rather than presented as a measurement made by this repository. Treat it as a property of the architecture, not as a benchmark of this package.
What the Vim class actually exposes
The package's public surface is narrow. The README shows one import, from vision_mamba import Vim, and one class. The constructor takes ten named arguments: dim, heads, dt_rank, dim_inner, d_state, num_classes, image_size, patch_size, channels, dropout and depth. Several names carry assumptions from the transformer world. dim is described in the README comment as the dimension of the transformer model and heads as the number of attention heads, even though the underlying mechanism is a state space model rather than attention. dt_rank is described as the rank of the dynamic routing matrix and d_state as the dimension of the state vector. Those two parameters are where the state space behaviour is configured, and they are the ones most likely to need tuning if you move away from the defaults shown. The README example sets dim=256, heads=8, dt_rank=32, dim_inner=256, d_state=256, num_classes=1000, image_size=224, patch_size=16, channels=3, dropout=0.1 and depth=12. That combination is a reasonable starting point for ImageNet-shaped input, but the README does not state which of these values come from the paper's configuration and which are arbitrary. If you need the paper's exact hyperparameters, check the paper, not this README.
The documented data flow: tensor in, tensor out
The README gives one worked example, and it is short. An input tensor is created with torch.randn(1, 3, 224, 224), which the comment identifies as shape (batch_size, channels, height, width). The model is instantiated with the arguments above, the tensor is passed to it, and the result is printed with out.shape and out. That is the whole documented contract: NCHW float input at the configured image_size and channels, and a tensor out. Because num_classes is 1000 in the example, the natural reading is that the output is a classification logit vector of length 1000, but the README never states the output shape explicitly. It prints it instead of describing it. The image_size and patch_size arguments imply an internal patching step, since 224 divided by 16 gives 14 patches per side, but how patches are ordered, whether positional information is added, and how the bidirectional state space scan is arranged across the patch sequence are not described anywhere in the supplied material. If you need to modify the internals, plan on reading the source rather than the documentation.
Getting it running: one install, one script
Installation is pip install vision-mamba. Note that the distribution name on PyPI uses a hyphen while the import name uses an underscore, which is why the README pairs the pip line with from vision_mamba import Vim. After that, the usage block is self-contained: import torch, import Vim, build the input tensor, construct the model with the keyword arguments shown, call it, print the shape. There are no config files, no environment variables, no CLI entry points and no YAML in the material provided. Every knob is a constructor keyword argument. That is convenient for embedding the model in an existing PyTorch project and inconvenient if you wanted a declarative configuration. The README also links a Discord server as the project homepage, which suggests questions are expected to be answered in chat rather than through documentation. For a package whose only documented behaviour is a forward pass, that is a thin support surface.
The training script is still a checkbox
The README ends with a Todo list containing two unchecked items: create training script for imagenet, and create a visual mamba for facial recognition. Neither is done as far as the supplied material shows. This is the single most important fact about the project for anyone evaluating it. You can instantiate the model. You cannot, from this repository alone, reproduce the results the paper reports, because the training procedure, the data pipeline, the optimiser settings, the augmentation policy and the schedule are not part of the package. A second unchecked item, the facial recognition variant, signals that the author intends the code to grow beyond the base architecture, which means the API may change. There are also no retrieved releases, so there is no tagged version to pin against. Combined with a last push date of 2026-08-29, the repository appears active but unpinned. Vendoring a specific commit is the only way to know what you are running.
Where a plain ViT or timm is the better choice
If your goal is to classify images and you do not specifically need the state space mechanism, a standard vision transformer from a mature model library is the lower-risk path. The difference in approach matters. A ViT applies self-attention across all patch pairs, which costs quadratic time in the number of patches; the Vision Mamba design instead runs a bidirectional state space scan over the patch sequence, which is the source of the efficiency claims in the paper. That is the reason to pick Vim. The reason not to pick this particular implementation is that a general model library typically ships training recipes, pretrained weights and a stable release history, and this repository ships none of those in the material provided. If you want the architecture for research, or you are building a custom pipeline where you control training end to end, the trade-off flips and a small readable implementation is an advantage. If you want a classifier working this week, it is not.
Licence and the cost of keeping up
The repository is MIT licensed, which permits commercial use, modification and redistribution provided the copyright notice and permission notice are retained. That is a permissive starting point, but the licence text itself is not reproduced in the README beyond the single word MIT, so confirm the LICENSE file exists in the tree before you rely on it. Maintenance cost is the harder question. The package has no releases, so upgrades are not versioned; pip install vision-mamba will pull whatever is current, and the constructor arguments you depend on could shift between installs. The README's Todo list indicates planned additions. The practical mitigation is to pin the dependency to a specific commit hash in your own requirements file and to re-read the Vim constructor signature whenever you bump it. There is no changelog in the supplied material, so the signature itself is the contract you have to check. Treat the ten keyword arguments as the stable surface and everything behind them as subject to change.
Editorial conclusion
Adopt this package if you want a pip-installable PyTorch module that instantiates the Vim architecture with a documented set of constructor arguments, and you are prepared to write your own training loop. Do not adopt it if you need a reproducible ImageNet training pipeline, pretrained weights, or published accuracy numbers: the README's own Todo list still has the ImageNet training script unchecked, and the repository ships no releases. Before wiring it into anything, verify three things against the repository at a specific commit: that the Vim constructor signature in the installed version matches the README example, that the output tensor shape is what your downstream code expects, and that the MIT licence file is present alongside the code you vendor.
Community notes