Below we describe how you can customize the Language Modeling Task. In our example, we add weight noise when training and freeze the backbone. For the purpose of this example we freeze the backbone within the Task, however this is recommended to be done via a Callback as seen in the Freeze Embeddings Callback.
Tasks are based of a AutoModel Transformer, which handles all the internal logic when running the forward pass through the model, and the loss calculation for a specific task. Below are the steps to customize a task within the LightningModule.
Inherit from Lightning Transformers Base Class
Add custom task logic
1. Inherit from Lightning Transformers Base Class¶
For our example, we inherit from the Language Modeling base class.
from lightning_transformers.task.nlp.language_modeling import LanguageModelingTransformer class MyLanguageModelingTransformer(LanguageModelingTransformer): ...
Typically you’d store the file within the
lightning_transformers/task/ directory, in the appropriate task folder.
In our example, we’d store our file in
2. Add Custom Task Logic¶
The class follows a standard
pl.LightningModule, thus all hooks and logic can be overridden easily.
Below we override the training step to add our logic, as well as
on_fit_start to freeze the model before training.
LMHeadAutoModel task provides separate keys for the backbone and the fully connected layer.
from lightning_transformers.task.nlp.language_modeling import LanguageModelingTransformer class MyLanguageModelingTransformer(LanguageModelingTransformer): def setup(self, stage): # Freeze BERT backbone for param in self.model.bert.parameters(): param.requires_grad = False def training_step(self, batch, batch_idx): loss = super().training_step(batch, batch_idx) # Add weight noise every training step with torch.no_grad(): for param in self.model.parameters(): param.add_(torch.randn(param.size()) * 0.1) return loss