Method get_optimizer_and_scheduler

textattack/trainer.py:327–386

Returns optimizer and scheduler to use for training. If you are overriding this method and do not want to use a scheduler, simply return :obj:`None` for scheduler. Args: model (:obj:`torch.nn.Module`): Model to be trained. Pass its parameters to o

(self, model, num_training_steps)

        wandb.log(log, step=step)

    def get_optimizer_and_scheduler(self, model, num_training_steps):
        """Returns optimizer and scheduler to use for training. If you are
        overriding this method and do not want to use a scheduler, simply
        return :obj:`None` for scheduler.

        Args:
            model (:obj:`torch.nn.Module`):
                Model to be trained. Pass its parameters to optimizer for training.
            num_training_steps (:obj:`int`):
                Number of total training steps.
        Returns:
            Tuple of optimizer and scheduler :obj:`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler._LRScheduler]`
        """
        if isinstance(model, torch.nn.DataParallel):
            model = model.module

        if isinstance(model, transformers.PreTrainedModel):
            # Reference https://huggingface.co/transformers/training.html
            param_optimizer = list(model.named_parameters())
            no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
            optimizer_grouped_parameters = [
                {
                    "params": [
                        p
                        for n, p in param_optimizer
                        if not any(nd in n for nd in no_decay)
                    ],
                    "weight_decay": self.training_args.weight_decay,
                },
                {
                    "params": [
                        p for n, p in param_optimizer if any(nd in n for nd in no_decay)
                    ],
                    "weight_decay": 0.0,
                },
            ]

            optimizer = torch.optim.AdamW(
                optimizer_grouped_parameters, lr=self.training_args.learning_rate
            )
            if isinstance(self.training_args.num_warmup_steps, float):
                num_warmup_steps = math.ceil(
                    self.training_args.num_warmup_steps * num_training_steps
                )
            else:
                num_warmup_steps = self.training_args.num_warmup_steps

            scheduler = transformers.optimization.get_linear_schedule_with_warmup(
                optimizer,
                num_warmup_steps=num_warmup_steps,
                num_training_steps=num_training_steps,
            )
        else:
            optimizer = torch.optim.Adam(
                filter(lambda x: x.requires_grad, model.parameters()),
                lr=self.training_args.learning_rate,
            )
            scheduler = None