jerilseb/mnist-classifier

1---2library_name: transformers3tags: []4---5 6# Usage7 8Define the model and config9```python10from transformers import PreTrainedModel, PretrainedConfig11import torch.nn as nn12import torch.nn.functional as F13 14class MNISTConfig(PretrainedConfig):15    model_type = "mnist_classifier"16 17    def __init__(self, input_size=784, hidden_size1=1024, hidden_size2=512, num_labels=10, **kwargs):18        super().__init__(**kwargs)19        self.input_size = input_size20        self.hidden_size1 = hidden_size121        self.hidden_size2 = hidden_size222        self.num_labels = num_labels23 24class MNISTClassifier(PreTrainedModel):25    config_class = MNISTConfig26 27    def __init__(self, config):28        super().__init__(config)29        self.layer1 = nn.Linear(config.input_size, config.hidden_size1)30        self.layer2 = nn.Linear(config.hidden_size1, config.hidden_size2)31        self.layer3 = nn.Linear(config.hidden_size2, config.num_labels)32 33    def forward(self, pixel_values):34        inputs = pixel_values.view(-1, self.config.input_size)35        outputs = self.layer1(inputs)36        outputs = F.leaky_relu(outputs)37        outputs = self.layer2(outputs)38        outputs = F.leaky_relu(outputs)39        outputs = self.layer3(outputs)40        return outputs41```42 43Register the model44 45```python46from transformers import AutoConfig, AutoModel47 48AutoConfig.register("mnist_classifier", MNISTConfig)49AutoModel.register(MNISTConfig, MNISTClassifier)50```51 52Run Inference53```python54from transformers import AutoConfig, AutoModel55import torch56 57config = AutoConfig.from_pretrained("jerilseb/mnist-classifier")58model = AutoModel.from_pretrained("jerilseb/mnist-classifier")59 60input_tensor = torch.randn(1, 28, 28)  # Single image, adjust batch size as needed61 62with torch.no_grad():63    output = model(input_tensor)64 65predicted_class = output.argmax(-1).item()66print(f"Predicted class: {predicted_class}")67```68