transformer模型转torchscript格式

from transformers import BertModel, BertTokenizer, BertConfig
import torch

enc = BertTokenizer.from_pretrained("bert-base-uncased")

# 输入文本tokenize
text = "[CLS] Who was Jim Henson ? [SEP] Jim Henson was a puppeteer [SEP]"
tokenized_text = enc.tokenize(text)

# 将一个token置为mask
masked_index = 8
tokenized_text[masked_index] = '[MASK]'
indexed_tokens = enc.convert_tokens_to_ids(tokenized_text)
segments_ids = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]

# 创建虚拟输入
tokens_tensor = torch.tensor([indexed_tokens])
segments_tensors = torch.tensor([segments_ids])
dummy_input = [tokens_tensor, segments_tensors]

# 初始化模型时将torchscript参数置为True
config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
    num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072, torchscript=True)

# 初始化模型
model = BertModel(config)

# 模型置为eval模式
model.eval()

# 也可以从pretrained初始化模型
model = BertModel.from_pretrained("bert-base-uncased", torchscript=True)

# 创建trace
traced_model = torch.jit.trace(model, [tokens_tensor, segments_tensors])
torch.jit.save(traced_model, "traced_bert.pt")

# 加载模型
loaded_model = torch.jit.load("traced_model.pt")
loaded_model.eval()

all_encoder_layers, pooled_output = loaded_model(dummy_input)

# 使用traced model进行推理
traced_model(tokens_tensor, segments_tensors)