"""
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
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)
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
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
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
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
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
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)
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}.")