File size: 1,885 Bytes
21bd514 58ed48f 21bd514 482076b 58ed48f 21bd514 58ed48f 21bd514 58ed48f 482076b 58ed48f 21bd514 58ed48f 21bd514 58ed48f 21bd514 58ed48f 21bd514 58ed48f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
---
library_name: transformers
tags: []
---
# Usage
Define the model and config
```python
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
```python
from transformers import AutoConfig, AutoModel
AutoConfig.register("mnist_classifier", MNISTConfig)
AutoModel.register(MNISTConfig, MNISTClassifier)
```
Run Inference
```python
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}")
```
|