def create_and_check_use_mems_train(
self,
config,
input_ids_1,
input_ids_2,
input_ids_q,
perm_mask,
input_mask,
target_mapping,
segment_ids,
lm_labels,
sequence_labels,
is_impossible_labels,
token_labels,
):
model = XLNetForSequenceClassification(config)
model.to(torch_device)
model.train()
train_size = input_ids_1.shape[0]
batch_size = 4
for i in range(train_size // batch_size + 1):
input_ids = input_ids_1[i:(i + 1) * batch_size]
labels = sequence_labels[i:(i + 1) * batch_size]
outputs = model(input_ids=input_ids,
labels=labels,
return_dict=True)
self.parent.assertIsNone(outputs.mems)
self.parent.assertIsNotNone(outputs.loss)
class XLNetSequenceClassifier:
"""XLNet-based sequence classifier"""
def __init__(
self,
language=Language.ENGLISHCASED,
num_labels=5,
cache_dir=".",
num_gpus=None,
num_epochs=1,
batch_size=8,
lr=5e-5,
adam_eps=1e-8,
warmup_steps=0,
weight_decay=0.0,
max_grad_norm=1.0,
):
"""Initializes the classifier and the underlying pretrained model.
Args:
language (Language, optional): The pretrained model's language.
Defaults to 'xlnet-base-cased'.
num_labels (int, optional): The number of unique labels in the
training data. Defaults to 5.
cache_dir (str, optional): Location of XLNet's cache directory.
Defaults to ".".
num_gpus (int, optional): The number of gpus to use.
If None is specified, all available GPUs
will be used. Defaults to None.
num_epochs (int, optional): Number of training epochs.
Defaults to 1.
batch_size (int, optional): Training batch size. Defaults to 8.
lr (float): Learning rate of the Adam optimizer. Defaults to 5e-5.
adam_eps (float, optional): term added to the denominator to improve
numerical stability. Defaults to 1e-8.
warmup_steps (int, optional): Number of steps in which to increase
learning rate linearly from 0 to 1. Defaults to 0.
weight_decay (float, optional): Weight decay. Defaults to 0.
max_grad_norm (float, optional): Maximum norm for the gradients. Defaults to 1.0
"""
if num_labels < 2:
raise ValueError("Number of labels should be at least 2.")
self.language = language
self.num_labels = num_labels
self.cache_dir = cache_dir
self.num_gpus = num_gpus
self.num_epochs = num_epochs
self.batch_size = batch_size
self.lr = lr
self.adam_eps = adam_eps
self.warmup_steps = warmup_steps
self.weight_decay = weight_decay
self.max_grad_norm = max_grad_norm
# create classifier
self.config = XLNetConfig.from_pretrained(
self.language.value, num_labels=num_labels, cache_dir=cache_dir
)
self.model = XLNetForSequenceClassification(self.config)
def fit(
self,
token_ids,
input_mask,
labels,
val_token_ids,
val_input_mask,
val_labels,
token_type_ids=None,
val_token_type_ids=None,
verbose=True,
logging_steps=0,
save_steps=0,
val_steps=0,
):
"""Fine-tunes the XLNet classifier using the given training data.
Args:
token_ids (list): List of training token id lists.
input_mask (list): List of input mask lists.
labels (list): List of training labels.
token_type_ids (list, optional): List of lists. Each sublist
contains segment ids indicating if the token belongs to
the first sentence(0) or second sentence(1). Only needed
for two-sentence tasks.
verbose (bool, optional): If True, shows the training progress and
loss values. Defaults to True.
"""
device, num_gpus = get_device(self.num_gpus)
self.model = move_to_device(self.model, device, self.num_gpus)
token_ids_tensor = torch.tensor(token_ids, dtype=torch.long)
input_mask_tensor = torch.tensor(input_mask, dtype=torch.long)
labels_tensor = torch.tensor(labels, dtype=torch.long)
val_token_ids_tensor = torch.tensor(val_token_ids, dtype=torch.long)
val_input_mask_tensor = torch.tensor(val_input_mask, dtype=torch.long)
val_labels_tensor = torch.tensor(val_labels, dtype=torch.long)
if token_type_ids:
token_type_ids_tensor = torch.tensor(token_type_ids, dtype=torch.long)
val_token_type_ids_tensor = torch.tensor(val_token_type_ids, dtype=torch.long)
train_dataset = TensorDataset(
token_ids_tensor, input_mask_tensor, token_type_ids_tensor, labels_tensor
)
val_dataset = TensorDataset(
val_token_ids_tensor,
val_input_mask_tensor,
val_token_type_ids_tensor,
val_labels_tensor,
)
else:
train_dataset = TensorDataset(token_ids_tensor, input_mask_tensor, labels_tensor)
val_dataset = TensorDataset(
val_token_ids_tensor, val_input_mask_tensor, val_labels_tensor
)
# define optimizer and model parameters
param_optimizer = list(self.model.named_parameters())
no_decay = ["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.weight_decay,
},
{
"params": [p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
"weight_decay": 0.0,
},
]
val_sampler = RandomSampler(val_dataset)
val_dataloader = DataLoader(val_dataset, sampler=val_sampler, batch_size=self.batch_size)
num_examples = len(token_ids)
num_batches = int(np.ceil(num_examples / self.batch_size))
num_train_optimization_steps = num_batches * self.num_epochs
optimizer = AdamW(optimizer_grouped_parameters, lr=self.lr, eps=self.adam_eps)
scheduler = WarmupLinearSchedule(
optimizer, warmup_steps=self.warmup_steps, t_total=num_train_optimization_steps
)
global_step = 0
self.model.train()
optimizer.zero_grad()
for epoch in range(self.num_epochs):
train_sampler = RandomSampler(train_dataset)
train_dataloader = DataLoader(
train_dataset, sampler=train_sampler, batch_size=self.batch_size
)
tr_loss = 0.0
logging_loss = 0.0
val_loss = 0.0
for i, batch in enumerate(tqdm(train_dataloader, desc="Iteration")):
if token_type_ids:
x_batch, mask_batch, token_type_ids_batch, y_batch = tuple(
t.to(device) for t in batch
)
else:
token_type_ids_batch = None
x_batch, mask_batch, y_batch = tuple(t.to(device) for t in batch)
outputs = self.model(
input_ids=x_batch,
token_type_ids=token_type_ids_batch,
attention_mask=mask_batch,
labels=y_batch,
)
loss = outputs[0] # model outputs are always tuple in pytorch-transformers
loss.sum().backward()
torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
tr_loss += loss.sum().item()
optimizer.step()
# Update learning rate schedule
scheduler.step()
optimizer.zero_grad()
global_step += 1
# logging of learning rate and loss
if logging_steps > 0 and global_step % logging_steps == 0:
mlflow.log_metric("learning rate", scheduler.get_lr()[0], step=global_step)
mlflow.log_metric(
"training loss",
(tr_loss - logging_loss) / (logging_steps * self.batch_size),
step=global_step,
)
logging_loss = tr_loss
# model checkpointing
if save_steps > 0 and global_step % save_steps == 0:
checkpoint_dir = os.path.join(os.getcwd(), "checkpoints")
if not os.path.isdir(checkpoint_dir):
os.makedirs(checkpoint_dir)
checkpoint_path = checkpoint_dir + "/" + str(global_step) + ".pth"
torch.save(self.model.state_dict(), checkpoint_path)
mlflow.log_artifact(checkpoint_path)
# model validation
if val_steps > 0 and global_step % val_steps == 0:
# run model on validation set
self.model.eval()
val_loss = 0.0
for j, val_batch in enumerate(val_dataloader):
if token_type_ids:
val_x_batch, val_mask_batch, val_token_type_ids_batch, val_y_batch = tuple(
t.to(device) for t in val_batch
)
else:
token_type_ids_batch = None
val_x_batch, val_mask_batch, val_y_batch = tuple(
t.to(device) for t in val_batch
)
val_outputs = self.model(
input_ids=val_x_batch,
token_type_ids=val_token_type_ids_batch,
attention_mask=val_mask_batch,
labels=val_y_batch,
)
vloss = val_outputs[0]
val_loss += vloss.sum().item()
mlflow.log_metric(
"validation loss", val_loss / len(val_dataset), step=global_step
)
self.model.train()
if verbose:
if i % ((num_batches // 10) + 1) == 0:
if val_loss > 0:
print(
"epoch:{}/{}; batch:{}->{}/{}; average training loss:{:.6f};\
average val loss:{:.6f}".format(
epoch + 1,
self.num_epochs,
i + 1,
min(i + 1 + num_batches // 10, num_batches),
num_batches,
tr_loss / (i + 1),
val_loss / (j + 1),
)
)
else:
print(
"epoch:{}/{}; batch:{}->{}/{}; average train loss:{:.6f}".format(
epoch + 1,
self.num_epochs,
i + 1,
min(i + 1 + num_batches // 10, num_batches),
num_batches,
tr_loss / (i + 1),
)
)
checkpoint_dir = os.path.join(os.getcwd(), "checkpoints")
if not os.path.isdir(checkpoint_dir):
os.makedirs(checkpoint_dir)
checkpoint_path = checkpoint_dir + "/" + "final" + ".pth"
torch.save(self.model.state_dict(), checkpoint_path)
mlflow.log_artifact(checkpoint_path)
# empty cache
del [x_batch, y_batch, mask_batch, token_type_ids_batch]
if val_steps > 0:
del [val_x_batch, val_y_batch, val_mask_batch, val_token_type_ids_batch]
torch.cuda.empty_cache()
def predict(
self,
token_ids,
input_mask,
token_type_ids=None,
num_gpus=None,
batch_size=8,
probabilities=False,
):
"""Scores the given dataset and returns the predicted classes.
Args:
token_ids (list): List of training token lists.
input_mask (list): List of input mask lists.
token_type_ids (list, optional): List of lists. Each sublist
contains segment ids indicating if the token belongs to
the first sentence(0) or second sentence(1). Only needed
for two-sentence tasks.
num_gpus (int, optional): The number of gpus to use.
If None is specified, all available GPUs
will be used. Defaults to None.
batch_size (int, optional): Scoring batch size. Defaults to 8.
probabilities (bool, optional):
If True, the predicted probability distribution
is also returned. Defaults to False.
Returns:
1darray, namedtuple(1darray, ndarray): Predicted classes or
(classes, probabilities) if probabilities is True.
"""
device, num_gpus = get_device(num_gpus)
self.model = move_to_device(self.model, device, num_gpus)
self.model.eval()
preds = []
with tqdm(total=len(token_ids)) as pbar:
for i in range(0, len(token_ids), batch_size):
start = i
end = start + batch_size
x_batch = torch.tensor(token_ids[start:end], dtype=torch.long, device=device)
mask_batch = torch.tensor(input_mask[start:end], dtype=torch.long, device=device)
token_type_ids_batch = torch.tensor(
token_type_ids[start:end], dtype=torch.long, device=device
)
with torch.no_grad():
pred_batch = self.model(
input_ids=x_batch,
token_type_ids=token_type_ids_batch,
attention_mask=mask_batch,
labels=None,
)
preds.append(pred_batch[0].cpu())
if i % batch_size == 0:
pbar.update(batch_size)
preds = np.concatenate(preds)
if probabilities:
return namedtuple("Predictions", "classes probabilities")(
preds.argmax(axis=1), nn.Softmax(dim=1)(torch.Tensor(preds)).numpy()
)
else:
return preds.argmax(axis=1)