File size: 3,730 Bytes
4943752 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
"""
PyTorch Model Wrapper
--------------------------
"""
import torch
from torch.nn import CrossEntropyLoss
import textattack
from .model_wrapper import ModelWrapper
torch.cuda.empty_cache()
class PyTorchModelWrapper(ModelWrapper):
"""Loads a PyTorch model (`nn.Module`) and tokenizer.
Args:
model (torch.nn.Module): PyTorch model
tokenizer: tokenizer whose output can be packed as a tensor and passed to the model.
No type requirement, but most have `tokenizer` method that accepts list of strings.
"""
def __init__(self, model, tokenizer):
if not isinstance(model, torch.nn.Module):
raise TypeError(
f"PyTorch model must be torch.nn.Module, got type {type(model)}"
)
self.model = model
self.tokenizer = tokenizer
def to(self, device):
self.model.to(device)
def __call__(self, text_input_list, batch_size=32):
model_device = next(self.model.parameters()).device
ids = self.tokenizer(text_input_list)
ids = torch.tensor(ids).to(model_device)
with torch.no_grad():
outputs = textattack.shared.utils.batch_model_predict(
self.model, ids, batch_size=batch_size
)
return outputs
def get_grad(self, text_input, loss_fn=CrossEntropyLoss()):
"""Get gradient of loss with respect to input tokens.
Args:
text_input (str): input string
loss_fn (torch.nn.Module): loss function. Default is `torch.nn.CrossEntropyLoss`
Returns:
Dict of ids, tokens, and gradient as numpy array.
"""
if not hasattr(self.model, "get_input_embeddings"):
raise AttributeError(
f"{type(self.model)} must have method `get_input_embeddings` that returns `torch.nn.Embedding` object that represents input embedding layer"
)
if not isinstance(loss_fn, torch.nn.Module):
raise ValueError("Loss function must be of type `torch.nn.Module`.")
self.model.train()
embedding_layer = self.model.get_input_embeddings()
original_state = embedding_layer.weight.requires_grad
embedding_layer.weight.requires_grad = True
emb_grads = []
def grad_hook(module, grad_in, grad_out):
emb_grads.append(grad_out[0])
emb_hook = embedding_layer.register_backward_hook(grad_hook)
self.model.zero_grad()
model_device = next(self.model.parameters()).device
ids = self.tokenizer([text_input])
ids = torch.tensor(ids).to(model_device)
predictions = self.model(ids)
output = predictions.argmax(dim=1)
loss = loss_fn(predictions, output)
loss.backward()
# grad w.r.t to word embeddings
# Fix for Issue #601
# Check if gradient has shape [max_sequence,1,_] ( when model input in transpose of input sequence)
if emb_grads[0].shape[1] == 1:
grad = torch.transpose(emb_grads[0], 0, 1)[0].cpu().numpy()
else:
# gradient has shape [1,max_sequence,_]
grad = emb_grads[0][0].cpu().numpy()
embedding_layer.weight.requires_grad = original_state
emb_hook.remove()
self.model.eval()
output = {"ids": ids[0].tolist(), "gradient": grad}
return output
def _tokenize(self, inputs):
"""Helper method that for `tokenize`
Args:
inputs (list[str]): list of input strings
Returns:
tokens (list[list[str]]): List of list of tokens as strings
"""
return [self.tokenizer.convert_ids_to_tokens(self.tokenizer(x)) for x in inputs]
|