Source code for textattack.trainer

"""
Trainer Class
=============
"""

import collections
import json
import logging
import math
import os

import scipy
import torch
import tqdm
import transformers

import textattack
from textattack.shared.utils import logger

from .attack import Attack
from .attack_args import AttackArgs
from .attack_results import MaximizedAttackResult, SuccessfulAttackResult
from .attacker import Attacker
from .model_args import HUGGINGFACE_MODELS
from .models.helpers import LSTMForClassification, WordCNNForClassification
from .models.wrappers import ModelWrapper
from .training_args import CommandLineTrainingArgs, TrainingArgs


[docs]class Trainer: """Trainer is training and eval loop for adversarial training. It is designed to work with PyTorch and Transformers models. Args: model_wrapper (:class:`~textattack.models.wrappers.ModelWrapper`): Model wrapper containing both the model and the tokenizer. task_type (:obj:`str`, `optional`, defaults to :obj:`"classification"`): The task that the model is trained to perform. Currently, :class:`~textattack.Trainer` supports two tasks: (1) :obj:`"classification"`, (2) :obj:`"regression"`. attack (:class:`~textattack.Attack`): :class:`~textattack.Attack` used to generate adversarial examples for training. train_dataset (:class:`~textattack.datasets.Dataset`): Dataset for training. eval_dataset (:class:`~textattack.datasets.Dataset`): Dataset for evaluation training_args (:class:`~textattack.TrainingArgs`): Arguments for training. Example:: >>> import textattack >>> import transformers >>> model = transformers.AutoModelForSequenceClassification.from_pretrained("bert-base-uncased") >>> tokenizer = transformers.AutoTokenizer.from_pretrained("bert-base-uncased") >>> model_wrapper = textattack.models.wrappers.HuggingFaceModelWrapper(model, tokenizer) >>> # We only use DeepWordBugGao2018 to demonstration purposes. >>> attack = textattack.attack_recipes.DeepWordBugGao2018.build(model_wrapper) >>> train_dataset = textattack.datasets.HuggingFaceDataset("imdb", split="train") >>> eval_dataset = textattack.datasets.HuggingFaceDataset("imdb", split="test") >>> # Train for 3 epochs with 1 initial clean epochs, 1000 adversarial examples per epoch, learning rate of 5e-5, and effective batch size of 32 (8x4). >>> training_args = textattack.TrainingArgs( ... num_epochs=3, ... num_clean_epochs=1, ... num_train_adv_examples=1000, ... learning_rate=5e-5, ... per_device_train_batch_size=8, ... gradient_accumulation_steps=4, ... log_to_tb=True, ... ) >>> trainer = textattack.Trainer( ... model_wrapper, ... "classification", ... attack, ... train_dataset, ... eval_dataset, ... training_args ... ) >>> trainer.train() .. note:: When using :class:`~textattack.Trainer` with `parallel=True` in :class:`~textattack.TrainingArgs`, make sure to protect the “entry point” of the program by using :obj:`if __name__ == '__main__':`. If not, each worker process used for generating adversarial examples will execute the training code again. """ def __init__( self, model_wrapper, task_type="classification", attack=None, train_dataset=None, eval_dataset=None, training_args=None, ): assert isinstance( model_wrapper, ModelWrapper ), f"`model_wrapper` must be of type `textattack.models.wrappers.ModelWrapper`, but got type `{type(model_wrapper)}`." # TODO: Support seq2seq training assert task_type in { "classification", "regression", }, '`task_type` must either be "classification" or "regression"' if attack: assert isinstance( attack, Attack ), f"`attack` argument must be of type `textattack.Attack`, but got type of `{type(attack)}`." if id(model_wrapper) != id(attack.goal_function.model): logger.warn( "`model_wrapper` and the victim model of `attack` are not the same model." ) if train_dataset: assert isinstance( train_dataset, textattack.datasets.Dataset ), f"`train_dataset` must be of type `textattack.datasets.Dataset`, but got type `{type(train_dataset)}`." if eval_dataset: assert isinstance( eval_dataset, textattack.datasets.Dataset ), f"`eval_dataset` must be of type `textattack.datasets.Dataset`, but got type `{type(eval_dataset)}`." if training_args: assert isinstance( training_args, TrainingArgs ), f"`training_args` must be of type `textattack.TrainingArgs`, but got type `{type(training_args)}`." else: training_args = TrainingArgs() if not hasattr(model_wrapper, "model"): raise ValueError("Cannot detect `model` in `model_wrapper`") else: assert isinstance( model_wrapper.model, torch.nn.Module ), f"`model` in `model_wrapper` must be of type `torch.nn.Module`, but got type `{type(model_wrapper.model)}`." if not hasattr(model_wrapper, "tokenizer"): raise ValueError("Cannot detect `tokenizer` in `model_wrapper`") self.model_wrapper = model_wrapper self.task_type = task_type self.attack = attack self.train_dataset = train_dataset self.eval_dataset = eval_dataset self.training_args = training_args self._metric_name = ( "pearson_correlation" if self.task_type == "regression" else "accuracy" ) if self.task_type == "regression": self.loss_fct = torch.nn.MSELoss(reduction="none") else: self.loss_fct = torch.nn.CrossEntropyLoss(reduction="none") self._global_step = 0 def _generate_adversarial_examples(self, epoch): """Generate adversarial examples using attacker.""" assert ( self.attack is not None ), "`attack` is `None` but attempting to generate adversarial examples." base_file_name = f"attack-train-{epoch}" log_file_name = os.path.join(self.training_args.output_dir, base_file_name) logger.info("Attacking model to generate new adversarial training set...") if isinstance(self.training_args.num_train_adv_examples, float): num_train_adv_examples = math.ceil( len(self.train_dataset) * self.training_args.num_train_adv_examples ) else: num_train_adv_examples = self.training_args.num_train_adv_examples # Use Different AttackArgs based on num_train_adv_examples value. # If num_train_adv_examples >= 0 , num_train_adv_examples is # set as number of successful examples. # If num_train_adv_examples == -1 , num_examples is set to -1 to # generate example for all of training data. if num_train_adv_examples >= 0: attack_args = AttackArgs( num_successful_examples=num_train_adv_examples, num_examples_offset=0, query_budget=self.training_args.query_budget_train, shuffle=True, parallel=self.training_args.parallel, num_workers_per_device=self.training_args.attack_num_workers_per_device, disable_stdout=True, silent=True, log_to_txt=log_file_name + ".txt", log_to_csv=log_file_name + ".csv", ) elif num_train_adv_examples == -1: # set num_examples when num_train_adv_examples = -1 attack_args = AttackArgs( num_examples=num_train_adv_examples, num_examples_offset=0, query_budget=self.training_args.query_budget_train, shuffle=True, parallel=self.training_args.parallel, num_workers_per_device=self.training_args.attack_num_workers_per_device, disable_stdout=True, silent=True, log_to_txt=log_file_name + ".txt", log_to_csv=log_file_name + ".csv", ) else: assert False, "num_train_adv_examples is negative and not equal to -1." attacker = Attacker(self.attack, self.train_dataset, attack_args=attack_args) results = attacker.attack_dataset() attack_types = collections.Counter(r.__class__.__name__ for r in results) total_attacks = ( attack_types["SuccessfulAttackResult"] + attack_types["FailedAttackResult"] ) success_rate = attack_types["SuccessfulAttackResult"] / total_attacks * 100 logger.info(f"Total number of attack results: {len(results)}") logger.info( f"Attack success rate: {success_rate:.2f}% [{attack_types['SuccessfulAttackResult']} / {total_attacks}]" ) # TODO: This will produce a bug if we need to manipulate ground truth output. # To Fix Issue #498 , We need to add the Non Output columns in one tuple to represent input columns # Since adversarial_example won't be an input to the model , we will have to remove it from the input # dictionary in collate_fn adversarial_examples = [ ( tuple(r.perturbed_result.attacked_text._text_input.values()) + ("adversarial_example",), r.perturbed_result.ground_truth_output, ) for r in results if isinstance(r, (SuccessfulAttackResult, MaximizedAttackResult)) ] # Name for column indicating if an example is adversarial is set as "_example_type". adversarial_dataset = textattack.datasets.Dataset( adversarial_examples, input_columns=self.train_dataset.input_columns + ("_example_type",), label_map=self.train_dataset.label_map, label_names=self.train_dataset.label_names, output_scale_factor=self.train_dataset.output_scale_factor, shuffle=False, ) return adversarial_dataset def _print_training_args( self, total_training_steps, train_batch_size, num_clean_epochs ): logger.info("***** Running training *****") logger.info(f" Num examples = {len(self.train_dataset)}") logger.info(f" Num epochs = {self.training_args.num_epochs}") logger.info(f" Num clean epochs = {num_clean_epochs}") logger.info( f" Instantaneous batch size per device = {self.training_args.per_device_train_batch_size}" ) logger.info( f" Total train batch size (w. parallel, distributed & accumulation) = {train_batch_size * self.training_args.gradient_accumulation_steps}" ) logger.info( f" Gradient accumulation steps = {self.training_args.gradient_accumulation_steps}" ) logger.info(f" Total optimization steps = {total_training_steps}") def _save_model_checkpoint( self, model, tokenizer, step=None, epoch=None, best=False, last=False ): # Save model checkpoint if step: dir_name = f"checkpoint-step-{step}" if epoch: dir_name = f"checkpoint-epoch-{epoch}" if best: dir_name = "best_model" if last: dir_name = "last_model" output_dir = os.path.join(self.training_args.output_dir, dir_name) if not os.path.exists(output_dir): os.makedirs(output_dir) if isinstance(model, torch.nn.DataParallel): model = model.module if isinstance(model, (WordCNNForClassification, LSTMForClassification)): model.save_pretrained(output_dir) elif isinstance(model, transformers.PreTrainedModel): model.save_pretrained(output_dir) tokenizer.save_pretrained(output_dir) else: state_dict = {k: v.cpu() for k, v in model.state_dict().items()} torch.save( state_dict, os.path.join(output_dir, "pytorch_model.bin"), ) def _tb_log(self, log, step): if not hasattr(self, "_tb_writer"): from torch.utils.tensorboard import SummaryWriter self._tb_writer = SummaryWriter(self.training_args.tb_log_dir) self._tb_writer.add_hparams(self.training_args.__dict__, {}) self._tb_writer.flush() for key in log: self._tb_writer.add_scalar(key, log[key], step) def _wandb_log(self, log, step): if not hasattr(self, "_wandb_init"): global wandb import wandb self._wandb_init = True wandb.init( project=self.training_args.wandb_project, config=self.training_args.__dict__, ) wandb.log(log, step=step)
[docs] 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 = transformers.optimization.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 return optimizer, scheduler
[docs] def get_train_dataloader(self, dataset, adv_dataset, batch_size): """Returns the :obj:`torch.utils.data.DataLoader` for training. Args: dataset (:class:`~textattack.datasets.Dataset`): Original training dataset. adv_dataset (:class:`~textattack.datasets.Dataset`): Adversarial examples generated from the original training dataset. :obj:`None` if no adversarial attack takes place. batch_size (:obj:`int`): Batch size for training. Returns: :obj:`torch.utils.data.DataLoader` """ # TODO: Add pairing option where we can pair original examples with adversarial examples. # Helper functions for collating data def collate_fn(data): input_texts = [] targets = [] is_adv_sample = [] for item in data: if "_example_type" in item[0].keys(): # Get example type value from OrderedDict and remove it adv = item[0].pop("_example_type") # with _example_type removed from item[0] OrderedDict # all other keys should be part of input _input, label = item if adv != "adversarial_example": raise ValueError( "`item` has length of 3 but last element is not for marking if the item is an `adversarial example`." ) else: is_adv_sample.append(True) else: # else `len(item)` is 2. _input, label = item is_adv_sample.append(False) if isinstance(_input, collections.OrderedDict): _input = tuple(_input.values()) else: _input = tuple(_input) if len(_input) == 1: _input = _input[0] input_texts.append(_input) targets.append(label) return input_texts, torch.tensor(targets), torch.tensor(is_adv_sample) if adv_dataset: dataset = torch.utils.data.ConcatDataset([dataset, adv_dataset]) train_dataloader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn, pin_memory=True, ) return train_dataloader
[docs] def get_eval_dataloader(self, dataset, batch_size): """Returns the :obj:`torch.utils.data.DataLoader` for evaluation. Args: dataset (:class:`~textattack.datasets.Dataset`): Dataset to use for evaluation. batch_size (:obj:`int`): Batch size for evaluation. Returns: :obj:`torch.utils.data.DataLoader` """ # Helper functions for collating data def collate_fn(data): input_texts = [] targets = [] for _input, label in data: if isinstance(_input, collections.OrderedDict): _input = tuple(_input.values()) else: _input = tuple(_input) if len(_input) == 1: _input = _input[0] input_texts.append(_input) targets.append(label) return input_texts, torch.tensor(targets) eval_dataloader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn, pin_memory=True, ) return eval_dataloader
[docs] def training_step(self, model, tokenizer, batch): """Perform a single training step on a batch of inputs. Args: model (:obj:`torch.nn.Module`): Model to train. tokenizer: Tokenizer used to tokenize input text. batch (:obj:`tuple[list[str], torch.Tensor, torch.Tensor]`): By default, this will be a tuple of input texts, targets, and boolean tensor indicating if the sample is an adversarial example. .. note:: If you override the :meth:`get_train_dataloader` method, then shape/type of :obj:`batch` will depend on how you created your batch. Returns: :obj:`tuple[torch.Tensor, torch.Tensor, torch.Tensor]` where - **loss**: :obj:`torch.FloatTensor` of shape 1 containing the loss. - **preds**: :obj:`torch.FloatTensor` of model's prediction for the batch. - **targets**: :obj:`torch.Tensor` of model's targets (e.g. labels, target values). """ input_texts, targets, is_adv_sample = batch _targets = targets targets = targets.to(textattack.shared.utils.device) if isinstance(model, transformers.PreTrainedModel) or ( isinstance(model, torch.nn.DataParallel) and isinstance(model.module, transformers.PreTrainedModel) ): input_ids = tokenizer( input_texts, padding="max_length", return_tensors="pt", truncation=True, ) input_ids.to(textattack.shared.utils.device) logits = model(**input_ids)[0] else: input_ids = tokenizer(input_texts) if not isinstance(input_ids, torch.Tensor): input_ids = torch.tensor(input_ids) input_ids = input_ids.to(textattack.shared.utils.device) logits = model(input_ids) if self.task_type == "regression": loss = self.loss_fct(logits.squeeze(), targets.squeeze()) preds = logits else: loss = self.loss_fct(logits, targets) preds = logits.argmax(dim=-1) sample_weights = torch.ones( is_adv_sample.size(), device=textattack.shared.utils.device ) sample_weights[is_adv_sample] *= self.training_args.alpha loss = loss * sample_weights loss = torch.mean(loss) preds = preds.cpu() return loss, preds, _targets
[docs] def evaluate_step(self, model, tokenizer, batch): """Perform a single evaluation step on a batch of inputs. Args: model (:obj:`torch.nn.Module`): Model to train. tokenizer: Tokenizer used to tokenize input text. batch (:obj:`tuple[list[str], torch.Tensor]`): By default, this will be a tuple of input texts and target tensors. .. note:: If you override the :meth:`get_eval_dataloader` method, then shape/type of :obj:`batch` will depend on how you created your batch. Returns: :obj:`tuple[torch.Tensor, torch.Tensor]` where - **preds**: :obj:`torch.FloatTensor` of model's prediction for the batch. - **targets**: :obj:`torch.Tensor` of model's targets (e.g. labels, target values). """ input_texts, targets = batch _targets = targets targets = targets.to(textattack.shared.utils.device) if isinstance(model, transformers.PreTrainedModel): input_ids = tokenizer( input_texts, padding="max_length", return_tensors="pt", truncation=True, ) input_ids.to(textattack.shared.utils.device) logits = model(**input_ids)[0] else: input_ids = tokenizer(input_texts) if not isinstance(input_ids, torch.Tensor): input_ids = torch.tensor(input_ids) input_ids = input_ids.to(textattack.shared.utils.device) logits = model(input_ids) if self.task_type == "regression": preds = logits else: preds = logits.argmax(dim=-1) return preds.cpu(), _targets
[docs] 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 ) # if num_train_adv_examples is int and >=0 then it is taken as value. elif ( isinstance(self.training_args.num_train_adv_examples, int) and self.training_args.num_train_adv_examples >= 0 ): total_adv_training_data_length = self.training_args.num_train_adv_examples # if num_train_adv_examples is = -1 , we generate all possible adv examples. # Max number of all possible adv examples would be equal to train_dataset. else: total_adv_training_data_length = len(self.train_dataset) # Based on total_adv_training_data_length calculation , find total total_adv_training_steps total_adv_training_steps = math.ceil( (len(self.train_dataset) + total_adv_training_data_length) / (train_batch_size * self.training_args.gradient_accumulation_steps) ) * (self.training_args.num_epochs - num_clean_epochs) total_training_steps = total_clean_training_steps + total_adv_training_steps optimizer, scheduler = self.get_optimizer_and_scheduler( model, total_training_steps ) self._print_training_args( total_training_steps, train_batch_size, num_clean_epochs ) model.to(textattack.shared.utils.device) # Variables across epochs self._total_loss = 0.0 self._current_loss = 0.0 self._last_log_step = 0 # `best_score` is used to keep track of the best model across training. # Could be loss, accuracy, or other metrics. best_eval_score = 0.0 best_eval_score_epoch = 0 best_model_path = None epochs_since_best_eval_score = 0 for epoch in range(1, self.training_args.num_epochs + 1): logger.info("==========================================================") logger.info(f"Epoch {epoch}") if self.attack and epoch > num_clean_epochs: if ( epoch - num_clean_epochs - 1 ) % self.training_args.attack_epoch_interval == 0: # only generate a new adversarial training set every self.training_args.attack_period epochs after the clean epochs # adv_dataset is instance of `textattack.datasets.Dataset` model.eval() adv_dataset = self._generate_adversarial_examples(epoch) model.train() model.to(textattack.shared.utils.device) else: adv_dataset = None else: logger.info(f"Running clean epoch {epoch}/{num_clean_epochs}") adv_dataset = None train_dataloader = self.get_train_dataloader( self.train_dataset, adv_dataset, train_batch_size ) model.train() # Epoch variables all_preds = [] all_targets = [] prog_bar = tqdm.tqdm( train_dataloader, desc="Iteration", position=0, leave=True, dynamic_ncols=True, ) for step, batch in enumerate(prog_bar): loss, preds, targets = self.training_step(model, tokenizer, batch) if isinstance(model, torch.nn.DataParallel): loss = loss.mean() loss = loss / self.training_args.gradient_accumulation_steps loss.backward() loss = loss.item() self._total_loss += loss self._current_loss += loss all_preds.append(preds) all_targets.append(targets) if (step + 1) % self.training_args.gradient_accumulation_steps == 0: optimizer.step() if scheduler: scheduler.step() optimizer.zero_grad() self._global_step += 1 if self._global_step > 0: prog_bar.set_description( f"Loss {self._total_loss/self._global_step:.5f}" ) # TODO: Better way to handle TB and Wandb logging if (self._global_step > 0) and ( self._global_step % self.training_args.logging_interval_step == 0 ): lr_to_log = ( scheduler.get_last_lr()[0] if scheduler else self.training_args.learning_rate ) if self._global_step - self._last_log_step >= 1: loss_to_log = round( self._current_loss / (self._global_step - self._last_log_step), 4, ) else: loss_to_log = round(self._current_loss, 4) log = {"train/loss": loss_to_log, "train/learning_rate": lr_to_log} if self.training_args.log_to_tb: self._tb_log(log, self._global_step) if self.training_args.log_to_wandb: self._wandb_log(log, self._global_step) self._current_loss = 0.0 self._last_log_step = self._global_step # Save model checkpoint to file. if self.training_args.checkpoint_interval_steps: if ( self._global_step > 0 and ( self._global_step % self.training_args.checkpoint_interval_steps ) == 0 ): self._save_model_checkpoint( model, tokenizer, step=self._global_step ) preds = torch.cat(all_preds) targets = torch.cat(all_targets) if self._metric_name == "accuracy": correct_predictions = (preds == targets).sum().item() accuracy = correct_predictions / len(targets) metric_log = {"train/train_accuracy": accuracy} logger.info(f"Train accuracy: {accuracy*100:.2f}%") else: pearson_correlation, pearson_pvalue = scipy.stats.pearsonr( preds, targets ) metric_log = { "train/pearson_correlation": pearson_correlation, "train/pearson_pvalue": pearson_pvalue, } logger.info(f"Train Pearson correlation: {pearson_correlation:.4f}%") if len(targets) > 0: if self.training_args.log_to_tb: self._tb_log(metric_log, epoch) if self.training_args.log_to_wandb: metric_log["epoch"] = epoch self._wandb_log(metric_log, self._global_step) # Evaluate after each epoch. eval_score = self.evaluate() if self.training_args.log_to_tb: self._tb_log({f"eval/{self._metric_name}": eval_score}, epoch) if self.training_args.log_to_wandb: self._wandb_log( {f"eval/{self._metric_name}": eval_score, "epoch": epoch}, self._global_step, ) if ( self.training_args.checkpoint_interval_epochs and (epoch % self.training_args.checkpoint_interval_epochs) == 0 ): self._save_model_checkpoint(model, tokenizer, epoch=epoch) if eval_score > best_eval_score: best_eval_score = eval_score best_eval_score_epoch = epoch epochs_since_best_eval_score = 0 self._save_model_checkpoint(model, tokenizer, best=True) logger.info( f"Best score found. Saved model to {self.training_args.output_dir}/best_model/" ) else: epochs_since_best_eval_score += 1 if self.training_args.early_stopping_epochs and ( epochs_since_best_eval_score > self.training_args.early_stopping_epochs ): logger.info( f"Stopping early since it's been {self.training_args.early_stopping_epochs} steps since validation score increased." ) break if self.training_args.log_to_tb: self._tb_writer.flush() # Finish training if isinstance(model, torch.nn.DataParallel): model = model.module if self.training_args.load_best_model_at_end: best_model_path = os.path.join(self.training_args.output_dir, "best_model") if hasattr(model, "from_pretrained"): model = model.__class__.from_pretrained(best_model_path) else: model = model.load_state_dict( torch.load(os.path.join(best_model_path, "pytorch_model.bin")) ) if self.training_args.save_last: self._save_model_checkpoint(model, tokenizer, last=True) self.model_wrapper.model = model self._write_readme(best_eval_score, best_eval_score_epoch, train_batch_size)
[docs] def evaluate(self): """Evaluate the model on given evaluation dataset.""" if not self.eval_dataset: raise ValueError("No `eval_dataset` available for training.") logging.info("Evaluating model on evaluation dataset.") model = self.model_wrapper.model tokenizer = self.model_wrapper.tokenizer model.eval() all_preds = [] all_targets = [] if isinstance(model, torch.nn.DataParallel): num_gpus = torch.cuda.device_count() eval_batch_size = self.training_args.per_device_eval_batch_size * num_gpus else: eval_batch_size = self.training_args.per_device_eval_batch_size eval_dataloader = self.get_eval_dataloader(self.eval_dataset, eval_batch_size) with torch.no_grad(): for step, batch in enumerate(eval_dataloader): preds, targets = self.evaluate_step(model, tokenizer, batch) all_preds.append(preds) all_targets.append(targets) preds = torch.cat(all_preds) targets = torch.cat(all_targets) if self.task_type == "regression": pearson_correlation, pearson_p_value = scipy.stats.pearsonr(preds, targets) eval_score = pearson_correlation else: correct_predictions = (preds == targets).sum().item() accuracy = correct_predictions / len(targets) eval_score = accuracy if self._metric_name == "accuracy": logger.info(f"Eval {self._metric_name}: {eval_score*100:.2f}%") else: logger.info(f"Eval {self._metric_name}: {eval_score:.4f}%") return eval_score
def _write_readme(self, best_eval_score, best_eval_score_epoch, train_batch_size): if isinstance(self.training_args, CommandLineTrainingArgs): model_name = self.training_args.model_name_or_path elif isinstance(self.model_wrapper.model, transformers.PreTrainedModel): if ( hasattr(self.model_wrapper.model.config, "_name_or_path") and self.model_wrapper.model.config._name_or_path in HUGGINGFACE_MODELS ): # TODO Better way than just checking HUGGINGFACE_MODELS ? model_name = self.model_wrapper.model.config._name_or_path elif hasattr(self.model_wrapper.model.config, "model_type"): model_name = self.model_wrapper.model.config.model_type else: model_name = "" else: model_name = "" if model_name: model_name = f"`{model_name}`" if ( isinstance(self.training_args, CommandLineTrainingArgs) and self.training_args.model_max_length ): model_max_length = self.training_args.model_max_length elif isinstance( self.model_wrapper.model, ( transformers.PreTrainedModel, LSTMForClassification, WordCNNForClassification, ), ): model_max_length = self.model_wrapper.tokenizer.model_max_length else: model_max_length = None if model_max_length: model_max_length_str = f" a maximum sequence length of {model_max_length}," else: model_max_length_str = "" if isinstance( self.train_dataset, textattack.datasets.HuggingFaceDataset ) and hasattr(self.train_dataset, "_name"): dataset_name = self.train_dataset._name if hasattr(self.train_dataset, "_subset"): dataset_name += f" ({self.train_dataset._subset})" elif isinstance( self.eval_dataset, textattack.datasets.HuggingFaceDataset ) and hasattr(self.eval_dataset, "_name"): dataset_name = self.eval_dataset._name if hasattr(self.eval_dataset, "_subset"): dataset_name += f" ({self.eval_dataset._subset})" else: dataset_name = None if dataset_name: dataset_str = ( "and the `{dataset_name}` dataset loaded using the `datasets` library" ) else: dataset_str = "" loss_func = ( "mean squared error" if self.task_type == "regression" else "cross-entropy" ) metric_name = ( "pearson correlation" if self.task_type == "regression" else "accuracy" ) epoch_info = f"{best_eval_score_epoch} epoch" + ( "s" if best_eval_score_epoch > 1 else "" ) readme_text = f""" ## TextAttack Model Card This {model_name} model was fine-tuned using TextAttack{dataset_str}. The model was fine-tuned for {self.training_args.num_epochs} epochs with a batch size of {train_batch_size}, {model_max_length_str} and an initial learning rate of {self.training_args.learning_rate}. Since this was a {self.task_type} task, the model was trained with a {loss_func} loss function. The best score the model achieved on this task was {best_eval_score}, as measured by the eval set {metric_name}, found after {epoch_info}. For more information, check out [TextAttack on Github](https://github.com/QData/TextAttack). """ readme_save_path = os.path.join(self.training_args.output_dir, "README.md") with open(readme_save_path, "w", encoding="utf-8") as f: f.write(readme_text.strip() + "\n") logger.info(f"Wrote README to {readme_save_path}.")