Method train

textattack/trainer.py:598–881

Train the model on given training dataset.

(self)

        return preds.cpu(), _targets

    def train(self):
        """Train the model on given training dataset."""
        if not self.train_dataset:
            raise ValueError("No `train_dataset` available for training.")

        textattack.shared.utils.set_seed(self.training_args.random_seed)
        if not os.path.exists(self.training_args.output_dir):
            os.makedirs(self.training_args.output_dir)

        # Save logger writes to file
        log_txt_path = os.path.join(self.training_args.output_dir, "train_log.txt")
        fh = logging.FileHandler(log_txt_path)
        fh.setLevel(logging.DEBUG)
        logger.addHandler(fh)
        logger.info(f"Writing logs to {log_txt_path}.")

        # Save original self.training_args to file
        args_save_path = os.path.join(
            self.training_args.output_dir, "training_args.json"
        )
        with open(args_save_path, "w", encoding="utf-8") as f:
            json.dump(self.training_args.__dict__, f)
        logger.info(f"Wrote original training args to {args_save_path}.")

        num_gpus = torch.cuda.device_count()
        tokenizer = self.model_wrapper.tokenizer
        model = self.model_wrapper.model

        if self.training_args.parallel and num_gpus > 1:
            # TODO: torch.nn.parallel.DistributedDataParallel
            # Supposedly faster than DataParallel, but requires more work to setup properly.
            model = torch.nn.DataParallel(model)
            logger.info(f"Training on {num_gpus} GPUs via `torch.nn.DataParallel`.")
            train_batch_size = self.training_args.per_device_train_batch_size * num_gpus
        else:
            train_batch_size = self.training_args.per_device_train_batch_size

        if self.attack is None:
            num_clean_epochs = self.training_args.num_epochs
        else:
            num_clean_epochs = self.training_args.num_clean_epochs

        total_clean_training_steps = (
            math.ceil(
                len(self.train_dataset)
                / (train_batch_size * self.training_args.gradient_accumulation_steps)
            )
            * num_clean_epochs
        )

        # calculate total_adv_training_data_length based on type of
        # num_train_adv_examples.
        # if num_train_adv_examples is float , num_train_adv_examples is a portion of train_dataset.
        if isinstance(self.training_args.num_train_adv_examples, float):
            total_adv_training_data_length = (
                len(self.train_dataset) * self.training_args.num_train_adv_examples
            )