Hi @gabeorlanski, I’m not aware of any dedicated tutorials for building custom models, but my suggestion would be to subclass PreTrainedModel
(check out how e.g. BertForSequenceClassification
is implemented) or one of the existing model classes. This has several advantages to using nn.Module
:
- You get all the helper methods like
from_pretrained
for free - Your custom model will play nice with the
Trainer
Depending on your use case, you can also override methods directly in the Trainer
- see here for a list of the available methods.