{"id":22354710,"url":"https://github.com/pprp/vision-mamba-cifar10","last_synced_at":"2025-06-17T01:33:53.776Z","repository":{"id":218302105,"uuid":"745560817","full_name":"pprp/Vision-Mamba-CIFAR10","owner":"pprp","description":null,"archived":false,"fork":false,"pushed_at":"2024-01-21T05:12:46.000Z","size":43,"stargazers_count":18,"open_issues_count":0,"forks_count":0,"subscribers_count":1,"default_branch":"main","last_synced_at":"2024-04-28T04:46:32.669Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/pprp.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null}},"created_at":"2024-01-19T15:50:54.000Z","updated_at":"2024-04-19T12:18:34.000Z","dependencies_parsed_at":"2024-01-21T04:45:08.118Z","dependency_job_id":null,"html_url":"https://github.com/pprp/Vision-Mamba-CIFAR10","commit_stats":null,"previous_names":["pprp/vision-mamba-cifar10"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/pprp%2FVision-Mamba-CIFAR10","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/pprp%2FVision-Mamba-CIFAR10/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/pprp%2FVision-Mamba-CIFAR10/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/pprp%2FVision-Mamba-CIFAR10/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/pprp","download_url":"https://codeload.github.com/pprp/Vision-Mamba-CIFAR10/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":228114855,"owners_count":17871742,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-12-04T13:14:46.513Z","updated_at":"2024-12-04T13:14:47.182Z","avatar_url":"https://github.com/pprp.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Vision-mamba-CIFAR\n\n## Introduction\n\n`Vision-mamba-CIFAR` is a toy project focused on image classification tasks using the CIFAR-10 dataset. Built with PyTorch and leveraging the state-of-the-art models from the `timm` library, this project aims to provide a high-level interface for training, evaluating, and deploying models trained on CIFAR-10. Whether you're a researcher, student, or hobbyist, `Vision-mamba-CIFAR` offers an accessible yet powerful way to jumpstart your projects in computer vision. Most of code is from `Vim`: https://github.com/hustvl/Vim \n\n## Features\n\n- Pre-trained model support from `timm`.\n- For the study of Mamba\n\n## Installation\n\n### Prerequisites\n\n- Python 3.8+\n- PyTorch 1.8+\n- torchvision\n- timm\n\n### Setup\n\nClone the repository to your local machine:\n\n```bash\ngit clone git@github.com:pprp/Vision-Mamba-CIFAR10.git\ncd Vision-mamba-CIFAR\n```\n\n### Installation \n\n```\npip install -r requirements\n```\n\n## Training \n\n```\nCUDA_VISIBLE_DEVICES=6 python main.py --data-set CIFAR \\\n    --model vim_tiny_patch16_224_bimambav2_final_pool_mean_abs_pos_embed_rope_also_residual --data-path ./data \\\n    --epochs 200 --lr 5e-4 \\\n    --batch-size 640 --gpu 6\n```\n\n```\nbash run.sh\n```\n\n\n## Simplified Code\n\n```python\nclass MambaBlock(nn.Module):\n    def __init__(self, args: ModelArgs):\n        \"\"\"A single Mamba block, as described in Figure 3 in Section 3.4 in the Mamba paper [1].\"\"\"\n        super().__init__()\n        self.args = args\n\n        self.in_proj = nn.Linear(args.d_model, args.d_inner * 2, bias=args.bias)\n\n        self.conv1d = nn.Conv1d(\n            in_channels=args.d_inner,\n            out_channels=args.d_inner,\n            bias=args.conv_bias,\n            kernel_size=args.d_conv,\n            groups=args.d_inner,\n            padding=args.d_conv - 1,\n        )\n\n        # x_proj takes in `x` and outputs the input-specific Δ, B, C\n        self.x_proj = nn.Linear(args.d_inner, args.dt_rank + args.d_state * 2, bias=False)\n        \n        # dt_proj projects Δ from dt_rank to d_in\n        self.dt_proj = nn.Linear(args.dt_rank, args.d_inner, bias=True)\n\n        A = repeat(torch.arange(1, args.d_state + 1), 'n -\u003e d n', d=args.d_inner)\n        self.A_log = nn.Parameter(torch.log(A))\n        self.D = nn.Parameter(torch.ones(args.d_inner))\n        self.out_proj = nn.Linear(args.d_inner, args.d_model, bias=args.bias)\n        \n\n    def forward(self, x):\n        \"\"\"Mamba block forward. This looks the same as Figure 3 in Section 3.4 in the Mamba paper [1].\n    \n        Args:\n            x: shape (b, l, d)    (See Glossary at top for definitions of b, l, d_in, n...)\n    \n        Returns:\n            output: shape (b, l, d)\n        \n        Official Implementation:\n            class Mamba, https://github.com/state-spaces/mamba/blob/main/mamba_ssm/modules/mamba_simple.py#L119\n            mamba_inner_ref(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/selective_scan_interface.py#L311\n            \n        \"\"\"\n        (b, l, d) = x.shape\n        \n        x_and_res = self.in_proj(x)  # shape (b, l, 2 * d_in)\n        (x, res) = x_and_res.split(split_size=[self.args.d_inner, self.args.d_inner], dim=-1)\n\n        x = rearrange(x, 'b l d_in -\u003e b d_in l')\n        x = self.conv1d(x)[:, :, :l]\n        x = rearrange(x, 'b d_in l -\u003e b l d_in')\n        \n        x = F.silu(x)\n\n        y = self.ssm(x)\n        \n        y = y * F.silu(res)\n        \n        output = self.out_proj(y)\n\n        return output\n\n    \n    def ssm(self, x):\n        \"\"\"Runs the SSM. See:\n            - Algorithm 2 in Section 3.2 in the Mamba paper [1]\n            - run_SSM(A, B, C, u) in The Annotated S4 [2]\n\n        Args:\n            x: shape (b, l, d_in)    (See Glossary at top for definitions of b, l, d_in, n...)\n    \n        Returns:\n            output: shape (b, l, d_in)\n\n        Official Implementation:\n            mamba_inner_ref(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/selective_scan_interface.py#L311\n            \n        \"\"\"\n        (d_in, n) = self.A_log.shape\n\n        # Compute ∆ A B C D, the state space parameters.\n        #     A, D are input independent (see Mamba paper [1] Section 3.5.2 \"Interpretation of A\" for why A isn't selective)\n        #     ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4,\n        #                                  and is why Mamba is called **selective** state spaces)\n        \n        A = -torch.exp(self.A_log.float())  # shape (d_in, n)\n        D = self.D.float()\n\n        x_dbl = self.x_proj(x)  # (b, l, dt_rank + 2*n)\n        \n        (delta, B, C) = x_dbl.split(split_size=[self.args.dt_rank, n, n], dim=-1)  # delta: (b, l, dt_rank). B, C: (b, l, n)\n        delta = F.softplus(self.dt_proj(delta))  # (b, l, d_in)\n        \n        y = self.selective_scan(x, delta, A, B, C, D)  # This is similar to run_SSM(A, B, C, u) in The Annotated S4 [2]\n        \n        return y\n\n    \n    def selective_scan(self, u, delta, A, B, C, D):\n        \"\"\"Does selective scan algorithm. See:\n            - Section 2 State Space Models in the Mamba paper [1]\n            - Algorithm 2 in Section 3.2 in the Mamba paper [1]\n            - run_SSM(A, B, C, u) in The Annotated S4 [2]\n\n        This is the classic discrete state space formula:\n            x(t + 1) = Ax(t) + Bu(t)\n            y(t)     = Cx(t) + Du(t)\n        except B and C (and the step size delta, which is used for discretization) are dependent on the input x(t).\n    \n        Args:\n            u: shape (b, l, d_in)    (See Glossary at top for definitions of b, l, d_in, n...)\n            delta: shape (b, l, d_in)\n            A: shape (d_in, n)\n            B: shape (b, l, n)\n            C: shape (b, l, n)\n            D: shape (d_in,)\n    \n        Returns:\n            output: shape (b, l, d_in)\n    \n        Official Implementation:\n            selective_scan_ref(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/selective_scan_interface.py#L86\n            Note: I refactored some parts out of `selective_scan_ref` out, so the functionality doesn't match exactly.\n            \n        \"\"\"\n        (b, l, d_in) = u.shape\n        n = A.shape[1]\n        \n        # Discretize continuous parameters (A, B)\n        # - A is discretized using zero-order hold (ZOH) discretization (see Section 2 Equation 4 in the Mamba paper [1])\n        # - B is discretized using a simplified Euler discretization instead of ZOH. From a discussion with authors:\n        #   \"A is the more important term and the performance doesn't change much with the simplification on B\"\n        deltaA = torch.exp(einsum(delta, A, 'b l d_in, d_in n -\u003e b l d_in n'))\n        deltaB_u = einsum(delta, B, u, 'b l d_in, b l n, b l d_in -\u003e b l d_in n')\n        \n        # Perform selective scan (see scan_SSM() in The Annotated S4 [2])\n        x = torch.zeros((b, d_in, n), device=deltaA.device)\n        ys = []    \n        for i in range(l):\n            x = deltaA[:, i] * x + deltaB_u[:, i]\n            y = einsum(x, C[:, i, :], 'b d_in n, b n -\u003e b d_in')\n            ys.append(y)\n        y = torch.stack(ys, dim=1)  # shape (b, l, d_in)\n        \n        y = y + u * D\n    \n        return y\n```","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fpprp%2Fvision-mamba-cifar10","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fpprp%2Fvision-mamba-cifar10","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fpprp%2Fvision-mamba-cifar10/lists"}