not-lain commited on
Commit
827dc1a
1 Parent(s): 34acb1a

Upload model

Browse files
Files changed (3) hide show
  1. MyModel.py +12 -5
  2. config.json +6 -1
  3. model.safetensors +1 -1
MyModel.py CHANGED
@@ -1,12 +1,15 @@
1
 
2
  from transformers import PreTrainedModel
3
- from .MyConfig import MnistConfig
4
  from torch import nn
5
  import torch.nn.functional as F
6
 
7
  class MnistModel(PreTrainedModel):
 
8
  config_class = MnistConfig
 
9
  def __init__(self, config):
 
10
  super().__init__(config)
11
  # use the config to instantiate our model
12
  self.conv1 = nn.Conv2d(1, config.conv1, kernel_size=5)
@@ -15,15 +18,19 @@ class MnistModel(PreTrainedModel):
15
  self.fc1 = nn.Linear(320, 50)
16
  self.fc2 = nn.Linear(50, 10)
17
  self.softmax = nn.Softmax(dim=-1)
 
18
  def forward(self, x,labels=None):
 
 
19
  x = F.relu(F.max_pool2d(self.conv1(x), 2))
20
  x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
21
  x = x.view(-1, 320)
22
  x = F.relu(self.fc1(x))
23
  x = F.dropout(x, training=self.training)
24
  x = self.fc2(x)
25
- output = self.softmax(x)
26
  if labels != None :
27
- print("continue training script here")
28
-
29
- return output
 
 
1
 
2
  from transformers import PreTrainedModel
3
+ from .MyConfig import MnistConfig # local import
4
  from torch import nn
5
  import torch.nn.functional as F
6
 
7
  class MnistModel(PreTrainedModel):
8
+ # pass the previously defined config class to the model
9
  config_class = MnistConfig
10
+
11
  def __init__(self, config):
12
+ # instantiate the model using the configuration
13
  super().__init__(config)
14
  # use the config to instantiate our model
15
  self.conv1 = nn.Conv2d(1, config.conv1, kernel_size=5)
 
18
  self.fc1 = nn.Linear(320, 50)
19
  self.fc2 = nn.Linear(50, 10)
20
  self.softmax = nn.Softmax(dim=-1)
21
+ self.criterion = nn.CrossEntropyLoss()
22
  def forward(self, x,labels=None):
23
+ # the labels parameter allows us to finetune our model
24
+ # with the Trainer API easily
25
  x = F.relu(F.max_pool2d(self.conv1(x), 2))
26
  x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
27
  x = x.view(-1, 320)
28
  x = F.relu(self.fc1(x))
29
  x = F.dropout(x, training=self.training)
30
  x = self.fc2(x)
31
+ logits = self.softmax(x)
32
  if labels != None :
33
+ # this will make your AI compatible with the trainer API
34
+ loss = self.criterion(logits, labels)
35
+ return {"loss": loss, "logits": logits}
36
+ return logits
config.json CHANGED
@@ -1,9 +1,14 @@
1
  {
 
 
 
2
  "auto_map": {
3
- "AutoConfig": "MyConfig.MnistConfig"
 
4
  },
5
  "conv1": 10,
6
  "conv2": 20,
7
  "model_type": "MobileNetV1",
 
8
  "transformers_version": "4.38.2"
9
  }
 
1
  {
2
+ "architectures": [
3
+ "MnistModel"
4
+ ],
5
  "auto_map": {
6
+ "AutoConfig": "MyConfig.MnistConfig",
7
+ "AutoModelForImageClassification": "MyModel.MnistModel"
8
  },
9
  "conv1": 10,
10
  "conv2": 20,
11
  "model_type": "MobileNetV1",
12
+ "torch_dtype": "float32",
13
  "transformers_version": "4.38.2"
14
  }
model.safetensors CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:6eee7d5198f4d57968bd06755727c4e89cd627b1b5fa7e0bfa42cf5a3bcd6697
3
  size 87976
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:7498727b72f4e20562223eb879e58b9410e06c37a678be0960dfb3650b11dd27
3
  size 87976