def benchmark_turbo_transformers(model_name: str, seq_len: int,
batch_size: int, n: int, enable_random: bool,
max_seq_len: int, min_seq_len: int,
num_threads: int, use_gpu: bool,
enable_mem_opt: bool):
import torch
import transformers
import turbo_transformers
import benchmark_helper
test_device = torch.device('cuda:0') if use_gpu else torch.device('cpu:0')
cfg = None
torch.set_grad_enabled(False)
if model_name == "bert":
cfg = transformers.BertConfig()
model = transformers.BertModel(cfg)
model.to(test_device)
model.eval()
model = turbo_transformers.BertModel.from_torch(model, backend="turbo")
elif model_name == "albert":
cfg = transformers.AlbertConfig(hidden_size=768,
num_attention_heads=12,
intermediate_size=3072)
model = transformers.AlbertModel(cfg)
model.to(test_device)
model.eval()
model = turbo_transformers.AlbertModel.from_torch(model)
elif model_name == "roberta":
cfg = transformers.RobertaConfig()
model = transformers.RobertaModel(cfg)
model.to(test_device)
model.eval()
model = turbo_transformers.RobertaModel.from_torch(model)
elif model_name == "distilbert":
cfg = transformers.DistilBertConfig()
model = transformers.DistilBertModel(cfg)
model.to(test_device)
model.eval()
model = turbo_transformers.DistilBertModel.from_torch(model)
else:
raise (f"benchmark does not support {model_name}")
turbo_transformers.set_num_threads(num_threads)
if enable_random:
if enable_mem_opt:
turbo_transformers.reset_allocator_schema("model-aware")
benchmark_helper.run_variable_model(model, use_gpu, n, max_seq_len,
min_seq_len, "turbo", num_threads,
cfg, enable_mem_opt, model_name)
if enable_mem_opt:
turbo_transformers.reset_allocator_schema("naive")
else:
input_ids = torch.randint(low=0,
high=cfg.vocab_size - 1,
size=(batch_size, seq_len),
dtype=torch.long,
device=test_device)
benchmark_helper.run_model(lambda: model(input_ids), use_gpu, n,
batch_size, seq_len, "turbo", num_threads,
enable_mem_opt, model_name)
def benchmark_torch(model_name: str, seq_len: int, batch_size: int, n: int,
enable_random: bool, max_seq_len: int, min_seq_len: int,
num_threads: int, use_gpu: bool, enable_mem_opt: bool):
import torch
import transformers
import benchmark_helper
test_device = torch.device('cuda:0') if use_gpu else torch.device('cpu:0')
torch.set_grad_enabled(False)
torch.set_num_threads(num_threads)
cfg = None
if model_name == "bert":
cfg = transformers.BertConfig()
model = transformers.BertModel(cfg)
elif model_name == "albert":
cfg = transformers.AlbertConfig()
model = transformers.AlbertModel(cfg)
elif model_name == "roberta":
cfg = transformers.RobertaConfig()
model = transformers.RobertaModel(cfg)
elif model_name == "distilbert":
cfg = transformers.DistilBertConfig()
model = transformers.DistilBertModel(cfg)
else:
raise (f"benchmark does not support {model_name}")
model.eval()
model.to(test_device)
# cfg = model.config # type: transformers.BertConfig
if enable_random:
benchmark_helper.run_variable_model(model, use_gpu, n, max_seq_len,
min_seq_len, "torch", num_threads,
cfg, enable_mem_opt, model_name)
else:
input_ids = torch.randint(low=0,
high=cfg.vocab_size - 1,
size=(batch_size, seq_len),
dtype=torch.long,
device=test_device)
benchmark_helper.run_model(lambda: model(input_ids), use_gpu, n,
batch_size, seq_len, "torch", num_threads,
enable_mem_opt, model_name)
plt.style.use('dark_background')
app = Flask(__name__)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = transformers.BertTokenizer.from_pretrained('./model/tokenizer')
PAD_INDEX = 0
BOS_INDEX = 101
EOS_INDEX = 102
MAX_IN_LEN = 512
MAX_OUT_LEN = 50
encoder_config = transformers.DistilBertConfig.from_pretrained('./model/encoder_config')
encoder = transformers.DistilBertModel(encoder_config)
bert_encoder = myBertModel(encoder, PAD_INDEX)
torch_decoder = myTorchDecoder(bert_encoder.get_embedding(), padding_idx=PAD_INDEX)
vocab_size_out, emb_size_decoder = torch_decoder.get_embedding_dim()
model = BertDecoderModel(bert_encoder, torch_decoder, emb_size_decoder, vocab_size_out).to(device)
model_dict = torch.load('./model/model.pt', map_location=device)
model_dict
model.load_state_dict(model_dict)
model.eval()
del encoder
del bert_encoder
del torch_decoder
def generate_onnx_model(model_name: str,
use_gpu: bool,
filename: str,
seq_len: int,
batch_size: int,
backend: str,
use_dynamic_axes: bool = False):
import transformers
import torch
test_device = torch.device(
'cuda:0') if backend == "GPU" and use_gpu else torch.device('cpu:0')
torch.set_grad_enabled(False)
if model_name == "bert":
# use a real model to check the correctness
if checkonnxrest:
model = transformers.BertModel.from_pretrained("bert-base-uncased")
else:
cfg = transformers.BertConfig()
model = transformers.BertModel(cfg)
elif model_name == "albert":
cfg = transformers.AlbertConfig()
model = transformers.AlbertModel(cfg)
elif model_name == "roberta":
cfg = transformers.RobertaConfig()
model = transformers.RobertaModel(cfg)
elif model_name == "distilbert":
cfg = transformers.DistilBertConfig()
model = transformers.DistilBertModel(cfg)
else:
raise (f"benchmark does not support {model_name}")
model.eval()
model.to(test_device)
cfg = model.config # type: transformers.BertConfig
input_ids = torch.randint(low=0,
high=cfg.vocab_size - 1,
size=(batch_size, seq_len),
dtype=torch.long,
device=test_device)
with open(filename, 'wb') as outf:
if not use_dynamic_axes:
torch.onnx.export(model=model, args=(input_ids, ), f=outf)
else:
torch.onnx.export(model=model,
args=(input_ids, ),
f=outf,
opset_version=11,
do_constant_folding=True,
input_names=['input'],
output_names=['output'],
dynamic_axes={
'input': [0, 1],
'output': [0, 1]
})
# If not intended to make onnxruntime support variable batch size and sequence length,
# you can unset the parameter `dynamic_axes`.
# For some model, you have to try `opset_version=12`
outf.flush()
return cfg.vocab_size, cfg