metadata
library_name: transformers
tags: []
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}")