CPU上での処理とInferentia上での処理との比較

import os
import time
import torch
import torch.neuron
from transformers import AutoTokenizer, AutoModelForSequenceClassification

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

# Load TorchScript back
model_cpu = torch.jit.load('bert_cpu.pt')
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")

# 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']

num_inferences = 1000

# warmup
paraphrase_classification_logits = model_cpu(*example_inputs_paraphrase)
paraphrase_classification_logits = model_neuron(*example_inputs_paraphrase)

# Run inference on CPUs, Display results
start = time.time()
for _ in range(num_inferences):
    paraphrase_classification_logits = model_cpu(*example_inputs_paraphrase)
elapsed_time = time.time() - start
print('By CPU         - num_inferences:{:>6}[sentences], elapsed_time:{:6.2f}[sec], Throughput:{:8.2f}[sentences/sec]'.format(num_inferences, elapsed_time, num_inferences / elapsed_time))

# Run inference on Neuron Cores, Display results
start = time.time()
for _ in range(num_inferences):
    paraphrase_classification_logits = model_neuron(*example_inputs_paraphrase)
elapsed_time = time.time() - start
print('By Neuron Core - num_inferences:{:>6}[sentences], elapsed_time:{:6.2f}[sec], Throughput:{:8.2f}[sentences/sec]'.format(num_inferences, elapsed_time, num_inferences / elapsed_time))

python infer_bert_perf.py

By CPU         - num_inferences:  1000[sentences], elapsed_time: 67.66[sec], Throughput:   14.78[sentences/sec]
By Neuron Core - num_inferences:  1000[sentences], elapsed_time:  7.20[sec], Throughput:  138.96[sentences/sec]

neuron-top

 NeuronCore Utilization
                  NC0                          NC1                         NC2                         NC3
 ND0  |||||||||||||||[94.57%] ||||||||||||||||||||[ 0.00%] ||||||||||||||||||||[ 0.00%] ||||||||||||||||||||[ 0.00%]

 vCPU and Memory Info
 System vCPU Usage ||||||||||||||||||||||||[31.60%, 4.19%]    Runtime vCPU Usage |||||||||||||||||||||||[23.98%, 0.38%]
 Runtime Memory Host ||||||||||||||||||[ 781.0KB/  15.2GB]    Runtime Memory Device  179.5MB

 Loaded Models
                                                                                                    Model ID
  [+] ND 0