Prediction MNIST-1D dataset with Mamba

import torch
import torch.nn as nn
import torch.nn.functional as F

from torchvision import transforms
from torchvision.datasets import MNIST


def create_dataloader(batch_size):
	data_train = torch.utils.data.DataLoader(
		MNIST(
			'~/mnist_data', train=True, download=True,
			transform=transforms.ToTensor(),
		),
		batch_size=batch_size,
		shuffle=True
	)
	
	data_test = torch.utils.data.DataLoader(
		MNIST(
			'~/mnist_data', train=False, download=True,
			transform=transforms.ToTensor(),
		),
		batch_size=batch_size,
		shuffle=True
	)
	
	return data_train, data_test


batch_size = 32
trainloader, testloader = create_dataloader(batch_size=batch_size)

import torch
from mamba_ssm import Mamba

  
class MambaMNISTClassifier(nn.Module):
	def __init__(
		self,
		length,
		dim,
		device="cuda"
	):
		super().__init__()
		self.mamba_model = Mamba(
			# This module uses roughly 3 * expand * d_model^2 parameters
			d_model=dim, # Model dimension d_model
			d_state=16, # SSM state expansion factor
			d_conv=4, # Local convolution width
			expand=2, # Block expansion factor
		).to(device)
		
		self.classifier = nn.Sequential(
			nn.Linear(length*dim, 10, ),
			nn.LogSoftmax(dim=1)
		).to(device)
		
	def forward(self, x):
		batch_size = x.shape[0]
		x = self.mamba_model(x)
		x = x.view(batch_size, -1)
		x = self.classifier(x)
		return x


length, dim = 784, 1
model = MambaMNISTClassifier(length, dim)

import torch
import torch.nn as nn
from tqdm.contrib import tenumerate

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

for epoch in range(10):
	print(f"Epoch: {epoch}")
	train_loss, test_loss = 0.0, 0.0
	model.train()

	for idx, samples in tenumerate(trainloader):
		data, label = samples
		# print(f"{data.shape=}, {label.shape=}") # (32, 1, 28, 28), (32,)

		inputs = data.view(batch_size, -1, 1).cuda()
		# print(f"{inputs.shape=}") # (32, 784, 1)
		targets = F.one_hot(label.view(batch_size), num_classes=10).float().cuda()

		optimizer.zero_grad()
		outputs = model(inputs)
		assert outputs.shape[1] == 10, f"{outputs.shape=}, {targets.shape=}"
		# print(f"{outputs.shape=}, {targets.shape=}") # (32, 10), (32, 10)

		loss = criterion(outputs, targets)
		loss.backward()
		optimizer.step()
		train_loss += loss.item()

	print("train loss: ", train_loss / len(trainloader))

	model.eval()
	with torch.no_grad():
		for idx, samples in enumerate(testloader):
			data, label = samples
			if idx == len(testloader) - 1:
				continue

			inputs = data.view(batch_size, -1, dim).cuda()
			targets = F.one_hot(label, num_classes=10).float().cuda()
			outputs = model(inputs)
			loss = criterion(outputs, targets)
			test_loss += loss.item()

	print("test loss: ", test_loss / len(testloader))