Upload model

config.json

@@ -0,0 +1,13 @@
+{
+  "architectures": [
+    "MyModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "test_model.MyModelConfig",
+    "AutoModel": "test_model.MyModel"
+  },
+  "input_dim": 10,
+  "layers_num": 3,
+  "torch_dtype": "float32",
+  "transformers_version": "4.24.0"
+}

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d3aaeadc87c87d041089e8ab0de07af2779cd0ff5952b2ecb477aa4301484f1f
+size 7047

test_model.py

@@ -0,0 +1,44 @@
+from transformers import PreTrainedModel
+from transformers import PretrainedConfig
+from typing import List
+import torch.nn as nn
+import torch
+
+
+class MyModelConfig(PretrainedConfig):
+
+    def __init__(# 每个参数都必须带有默认值，否则会报错
+        self,
+        input_dim=100,
+        layers_num=5,
+        **kwargs,
+    ):
+        self.input_dim = input_dim
+        self.layers_num = layers_num
+        super().__init__(**kwargs)
+
+class MyModel(PreTrainedModel):
+    config_class = MyModelConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        modules = []
+        assert config.layers_num >= 1
+        if config.layers_num == 1:
+            modules.append(nn.Linear(config.input_dim,1)) 
+        else:
+            modules.append(nn.Linear(config.input_dim,30))
+            for i in range(config.layers_num-2):
+                modules.append(nn.Linear(30,30))
+            modules.append(nn.Linear(30,1))
+        self.model = nn.ModuleList(modules)
+
+
+    def forward(self, tensor):
+        return self.model(tensor)
+
+if __name__ == '__main__':
+    save_config = MyModelConfig(input_dim=10,layers_num=3)
+    save_config.save_pretrained("custom-mymodel")
+    mymodel = MyModel(save_config)
+    torch.save(mymodel.state_dict(),'pytorch_model.bin') # 通常以此命名