jerilseb/mnist-classifier

Usage

Define the model and config

from transformers import PreTrainedModel, PretrainedConfig
import torch.nn as nn
import torch.nn.functional as F

class MNISTConfig(PretrainedConfig):
    model_type = "mnist_classifier"

    def __init__(self, input_size=784, hidden_size1=1024, hidden_size2=512, num_labels=10, **kwargs):
        super().__init__(**kwargs)
        self.input_size = input_size
        self.hidden_size1 = hidden_size1
        self.hidden_size2 = hidden_size2
        self.num_labels = num_labels

class MNISTClassifier(PreTrainedModel):
    config_class = MNISTConfig

    def __init__(self, config):
        super().__init__(config)
        self.layer1 = nn.Linear(config.input_size, config.hidden_size1)
        self.layer2 = nn.Linear(config.hidden_size1, config.hidden_size2)
        self.layer3 = nn.Linear(config.hidden_size2, config.num_labels)

    def forward(self, pixel_values):
        inputs = pixel_values.view(-1, self.config.input_size)
        outputs = self.layer1(inputs)
        outputs = F.leaky_relu(outputs)
        outputs = self.layer2(outputs)
        outputs = F.leaky_relu(outputs)
        outputs = self.layer3(outputs)
        return outputs

Register the model

from transformers import AutoConfig, AutoModel

AutoConfig.register("mnist_classifier", MNISTConfig)
AutoModel.register(MNISTConfig, MNISTClassifier)

Run Inference

from transformers import AutoConfig, AutoModel
import torch

config = AutoConfig.from_pretrained("jerilseb/mnist-classifier")
model = AutoModel.from_pretrained("jerilseb/mnist-classifier")

input_tensor = torch.randn(1, 28, 28)  # Single image, adjust batch size as needed

with torch.no_grad():
    output = model(input_tensor)

predicted_class = output.argmax(-1).item()
print(f"Predicted class: {predicted_class}")