Prepare optimizer and schedule (linear warmup and decay)
(self)
| 99 | self.model = self.model_type.from_pretrained(*args, **kwargs) |
| 100 | |
| 101 | def configure_optimizers(self): |
| 102 | "Prepare optimizer and schedule (linear warmup and decay)" |
| 103 | model = self.model |
| 104 | no_decay = ["bias", "LayerNorm.weight"] |
| 105 | optimizer_grouped_parameters = [ |
| 106 | { |
| 107 | "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], |
| 108 | "weight_decay": self.hparams.weight_decay, |
| 109 | }, |
| 110 | { |
| 111 | "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], |
| 112 | "weight_decay": 0.0, |
| 113 | }, |
| 114 | ] |
| 115 | optimizer = AdamW(optimizer_grouped_parameters, lr=self.hparams.learning_rate, eps=self.hparams.adam_epsilon) |
| 116 | self.opt = optimizer |
| 117 | return [optimizer] |
| 118 | |
| 119 | def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx, second_order_closure=None): |
| 120 | if self.trainer.use_tpu: |