コンパイルとInferentia上での実行

import torch.neuron

# Run torch.jit.trace to generate a TorchScript for running on CPU
model_cpu = torch.jit.trace(model, example_inputs_paraphrase)

# Run torch.neuron.trace to generate a TorchScript that is optimized by AWS Neuron
model_neuron = torch.neuron.trace(model, example_inputs_paraphrase)

# Save the TorchScripts for later use
model_cpu.save('bert_cpu.pt')
model_neuron.save('bert_neuron.pt')

time python compile_bert.py  

..
Compiler status PASS
INFO:Neuron:Number of arithmetic operators (post-compilation) before = 565, compiled = 548, percent compiled = 96.99%
INFO:Neuron:The neuron partitioner created 1 sub-graphs
INFO:Neuron:Neuron successfully compiled 1 sub-graphs, Total fused subgraphs = 1, Percent of model sub-graphs successfully compiled = 100.0%
INFO:Neuron:Compiled these operators (and operator counts) to Neuron:
INFO:Neuron: => aten::Int: 96
INFO:Neuron: => aten::add: 36
INFO:Neuron: => aten::contiguous: 12
INFO:Neuron: => aten::div: 12
INFO:Neuron: => aten::dropout: 38
INFO:Neuron: => aten::gelu: 12
INFO:Neuron: => aten::layer_norm: 25
INFO:Neuron: => aten::linear: 74
INFO:Neuron: => aten::matmul: 24
INFO:Neuron: => aten::permute: 48
INFO:Neuron: => aten::select: 1
INFO:Neuron: => aten::size: 96
INFO:Neuron: => aten::slice: 1
INFO:Neuron: => aten::softmax: 12
INFO:Neuron: => aten::tanh: 1
INFO:Neuron: => aten::transpose: 12
INFO:Neuron: => aten::view: 48
INFO:Neuron:Not compiled operators (and operator counts) to Neuron:
INFO:Neuron: => aten::Int: 1 [supported]
INFO:Neuron: => aten::add: 3 [supported]
INFO:Neuron: => aten::embedding: 3 [not supported]
INFO:Neuron: => aten::mul: 1 [supported]
INFO:Neuron: => aten::rsub: 1 [supported]
INFO:Neuron: => aten::size: 1 [supported]
INFO:Neuron: => aten::slice: 4 [supported]
INFO:Neuron: => aten::to: 1 [supported]
INFO:Neuron: => aten::unsqueeze: 2 [supported]

real    2m8.977s
user    2m5.834s
sys     0m9.073s

import torch
import torch.neuron
from transformers import AutoTokenizer

# Build tokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased-finetuned-mrpc")

# Load TorchScript back
model_neuron = torch.jit.load('bert_neuron.pt')

# Setup some example inputs
sequence_0 = "The company HuggingFace is based in New York City"
sequence_1 = "Apples are especially bad for your health"
sequence_2 = "HuggingFace's headquarters are situated in Manhattan"

paraphrase = tokenizer.encode_plus(sequence_0, sequence_2, max_length=128, padding='max_length', truncation=True, return_tensors="pt")
not_paraphrase = tokenizer.encode_plus(sequence_0, sequence_1, max_length=128, padding='max_length', truncation=True, return_tensors="pt")

# Convert example inputs to a format that is compatible with TorchScript tracing
example_inputs_paraphrase = paraphrase['input_ids'], paraphrase['attention_mask'], paraphrase['token_type_ids']
example_inputs_not_paraphrase = not_paraphrase['input_ids'], not_paraphrase['attention_mask'], not_paraphrase['token_type_ids']

# Verify the TorchScript works on both example inputs
paraphrase_classification_logits = model_neuron(*example_inputs_paraphrase)
not_paraphrase_classification_logits = model_neuron(*example_inputs_not_paraphrase)

classes = ['not paraphrase', 'paraphrase']
paraphrase_prediction = paraphrase_classification_logits[0][0].argmax().item()
not_paraphrase_prediction = not_paraphrase_classification_logits[0][0].argmax().item()
print('BERT says that "{}" and "{}" are {}'.format(sequence_0, sequence_2, classes[paraphrase_prediction]))
print('BERT says that "{}" and "{}" are {}'.format(sequence_0, sequence_1, classes[not_paraphrase_prediction]))

python infer_bert_neuron.py

BERT says that "The company HuggingFace is based in New York City" and "HuggingFace's headquarters are situated in Manhattan" are paraphrase
BERT says that "The company HuggingFace is based in New York City" and "Apples are especially bad for your health" are not paraphrase