将基于 Bert 的 PyTorch 模型导出到 CoreML。如何使 CoreML 模型适用于任何输入？

Question

我使用下面的代码将基于 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";

---

环境：

• 导出脚本：在 Ubuntu 20.04 上测试了 Python 3.10 和 torch 2.3.1（在 Windows 10 上不工作）。
• 预测脚本：必须在 macOS 10.13+ 上运行，因为 CoreML 模型仅支持 macOS 10.13+ 上的预测。

Answer 1

问题是您正在使用 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()

函数添加相同的参数。