Resources for using custom models with trainer

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.

1 Like