|
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') |
