我使用下面的代码将基于 Bert 的 PyTorch 模型导出到 CoreML。
自从用了
dummy_input = tokenizer("A French fan", return_tensors="pt")
CoreML 模型仅在 macOS 上测试时适用于该输入。如何使 CoreML 模型适用于任何输入(即任何文本)?
导出脚本:
# -*- coding: utf-8 -*-
"""Core ML Export
pip install transformers torch coremltools nltk
"""
import os
from transformers import AutoModelForTokenClassification, AutoTokenizer
import torch
import torch.nn as nn
import nltk
import coremltools as ct
nltk.download('punkt')
# Load the model and tokenizer
model_path = os.path.join('model')
model = AutoModelForTokenClassification.from_pretrained(model_path, local_files_only=True)
tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
# Modify the model's forward method to return a tuple
class ModifiedModel(nn.Module):
def __init__(self, model):
super(ModifiedModel, self).__init__()
self.model = model
self.device = model.device # Add the device attribute
def forward(self, input_ids, attention_mask, token_type_ids=None):
outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
return outputs.logits
modified_model = ModifiedModel(model)
# Export to Core ML
def convert_to_coreml(model, tokenizer):
# Define a dummy input for tracing
dummy_input = tokenizer("A French fan", return_tensors="pt")
dummy_input = {k: v.to(model.device) for k, v in dummy_input.items()}
# Trace the model with the dummy input
traced_model = torch.jit.trace(model, (
dummy_input['input_ids'], dummy_input['attention_mask'], dummy_input.get('token_type_ids')))
# Convert to Core ML
inputs = [
ct.TensorType(name="input_ids", shape=dummy_input['input_ids'].shape),
ct.TensorType(name="attention_mask", shape=dummy_input['attention_mask'].shape)
]
if 'token_type_ids' in dummy_input:
inputs.append(ct.TensorType(name="token_type_ids", shape=dummy_input['token_type_ids'].shape))
mlmodel = ct.convert(traced_model, inputs=inputs)
# Save the Core ML model
mlmodel.save("model.mlmodel")
print("Model exported to Core ML successfully")
convert_to_coreml(modified_model, tokenizer)
要使用导出的模型:
import os
from transformers import AutoModelForTokenClassification, AutoTokenizer
import torch
import torch.nn as nn
import nltk
import coremltools as ct
from coremltools.models import MLModel
import numpy as np
from transformers import AutoTokenizer
import nltk
nltk.download('punkt')
# Load the Core ML model
model = MLModel('model.mlmodel')
# Load the tokenizer
model_path = 'model'
tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
def prepare_input(text, tokenizer):
tokens = nltk.tokenize.word_tokenize(text)
tokenized_inputs = tokenizer(tokens, is_split_into_words=True, return_tensors="np")
input_ids = tokenized_inputs['input_ids'].astype(np.int32)
attention_mask = tokenized_inputs['attention_mask'].astype(np.int32)
input_data = {
'input_ids': input_ids,
'attention_mask': attention_mask
}
if 'token_type_ids' in tokenized_inputs:
input_data['token_type_ids'] = tokenized_inputs['token_type_ids'].astype(np.int32)
return input_data, tokens
def predict(text):
# Prepare the input
input_data, tokens = prepare_input(text, tokenizer)
# Make the prediction
prediction = model.predict(input_data)
# Extract the predicted labels
logits = prediction['output'] # Adjust this key according to your model's output
predicted_label = np.argmax(logits, axis=-1)[0]
# Display the results
for word, label in zip(tokens, predicted_label):
print(f"{word}: {model.model_description.outputDescriptions[0].dictionaryType.int64KeyType.stringDictionary[label]}")
# Test the model with a sentence
predict("A French fan")
该脚本仅适用于示例“法国粉丝”。当我尝试另一个示例
predict("A footbal fan is standing in the stadium.")
时,它触发了一个错误:
NSLocalizedDescription = "MultiArray shape (1 x 12) does not match the shape (1 x 5) specified in the model description";
环境:
问题是您正在使用 dummy_input 的形状(即 1 x 12)转换模型,并且
ct.convert()
函数假设这始终是您输入的形状。如果您希望能够使用任意数量的标记,直至模型的上下文长度,您需要将 shape 参数修改为如下所示:
inputs=[
ct.TensorType(name="input_ids", shape=(1, ct.RangeDim(1, 512)), dtype=np.int32),
ct.TensorType(name="attention_mask", shape=(1, ct.RangeDim(1, 512)), dtype=np.int32)
],
其中
1
中的 ct.RangeDim()
参数是最小上下文大小,512
参数是模型的最大上下文大小。
您还需要向
tokenizer()
函数添加一些参数,以确保它填充这些张量的其余部分,使其适合您的模型。您需要找到模型输入尺寸的尺寸,以便为其指定正确的尺寸。
dummy_input = tokenizer("This is a sample input", return_tensors="pt", padding='max_length', max_length=512)
在推理时,您也需要向
tokenizer()
函数添加相同的参数。