rolson24/Mamba-YOLO

the official pytorch implementation of “Mamba-YOLO：SSMs-based for Object Detection”

Mamba-YOLO

Official pytorch implementation of “Mamba-YOLO：SSMs-based for Object Detection”

Installation

# pip install required packages
conda create -n mambayolo -y python=3.11
conda activate mambayolo
pip3 install torch===2.3.0 torchvision torchaudio
pip install seaborn thop timm einops
cd selective_scan && pip install . && cd ..
pip install -v -e .

Training

python mbyolo_train.py --task train --data ultralytics/cfg/datasets/coco.yaml \
 --config ultralytics/cfg/models/v8/mamba-yolo.yaml \
--amp  --project ./output_dir/mscoco --name mambayolo_n

Acknowledgement

This repo is modified from open source real-time object detection codebase Ultralytics. The selective-scan from VMamba.