Huggingface PretrainedModel

import torch.nn as nn
from transformers import PreTrainedModel, PretrainedConfig
from transformers import AutoModel, AutoConfig

class MyConfig(PretrainedConfig):
    model_type = 'mymodel'
    def __init__(self, important_param=42, **kwargs):
        super().__init__(**kwargs)
        self.important_param = important_param

class MyModel(PreTrainedModel):
    config_class = MyConfig
    def __init__(self, config):
        super().__init__(config)
        self.config = config
        self.model = nn.Sequential(
                          nn.Linear(3, self.config.important_param),
                          nn.Sigmoid(),
                          nn.Linear(self.config.important_param, 1),
                          nn.Sigmoid()
                          )
    def forward(self, input):
        return self.model(input) 

config = MyConfig(4)
model = MyModel(config)
model.save_pretrained('./my_model_dir')

new_model = MyModel.from_pretrained('./my_model_dir')
new_model

AutoConfig.register("mymodel", MyConfig)
AutoModel.register(MyConfig, MyModel)

new_model = AutoModel.from_pretrained('./my_model_dir')
new_model

How to convert a PyTorch nn.Module into a HuggingFace PreTrainedModel object?

Given a simple neural net in Pytorch like: import torch.nn as nn net = nn.Sequential( nn.Linear(3, 4), nn.Sigmoid(), nn.Linear(4, 1), nn.Sigmoid() ).to(device) How d...

https://stackoverflow.com/questions/73948214/how-to-convert-a-pytorch-nn-module-into-a-huggingface-pretrainedmodel-object