本地运行训练好的NER模型时TFSMLayer输入类型不匹配报错

在Google Colab上训练完NER模型后，下载到Windows本地运行推理时遇到类型不匹配错误。Colab中直接用load_model加载模型可正常运行，但本地使用TFSMLayer加载时出现如下报错：

tensorflow.python.framework.errors_impl.InvalidArgumentError: Exception encountered when calling TFSMLayer.call().

cannot compute __inference_signature_wrapper_1285 as input #0(zero-based) was expected to be a int64 tensor but is a float tensor [Op:__inference_signature_wrapper_1285]

模型期望输入为int64类型张量，但实际传入TFSMLayer.call()的是float32类型，尝试用tf.cast转换后错误依旧，不清楚输入类型为何被自动转为float32，需要调试和解决建议。

import re
import pickle
import keras
import tensorflow as tf
from keras.models import Sequential
from keras.layers import TFSMLayer
import numpy as np

class CustomNonPaddingTokenLoss(keras.losses.Loss):
    def __init__(self, reduction=tf.keras.losses.Reduction.AUTO, name="custom_ner_loss"):
        super().__init__(reduction=reduction, name=name)

    def call(self, y_true, y_pred):
        loss_fn = keras.losses.SparseCategoricalCrossentropy(
            from_logits=False, reduction=self.reduction
        )
        loss = loss_fn(y_true, y_pred)
        mask = tf.cast((y_true > 0), dtype=tf.float32)
        loss = loss * mask
        return tf.reduce_sum(loss) / tf.reduce_sum(mask)

def map_record_to_training_data(record):
    record = tf.strings.split(record, sep="\t")
    length = tf.strings.to_number(record[0], out_type=tf.int32)
    tokens = record[1 : length + 1]
    tags = record[length + 1 :]
    tags = tf.strings.to_number(tags, out_type=tf.int64)
    tags += 1
    return tokens, tags

def lookup(tokens):
    # Load the list from the file
    with open('./resources/vocabulary.pkl', 'rb') as f:
        loaded_list = pickle.load(f)
    # The StringLookup class will convert tokens to token IDs
    lookup_layer = keras.layers.StringLookup(vocabulary=loaded_list)

    # No need to lowercase Vietnamese characters
    return lookup_layer(tokens)

def format_datatype(data):
    tokens =  [re.sub(r'[;,]', '', d) for d in data.split(' ')]
    #default is 0, since is for prediction
    ner_tags = [0 for d in data.split(' ')]

    #tab to separate
    string_input = str(len(tokens))+ "\t"+ "\t".join(tokens)+ "\t"+ "\t".join(map(str, ner_tags))
    string_input = tf.data.Dataset.from_tensor_slices([string_input])

    
    finalize_input = (string_input.map(map_record_to_training_data)
                      .map(lambda x, y: (lookup(x),  y))
                      .padded_batch(1)
                      )

    return finalize_input

def prediction(data):
    # Register the custom loss function with TensorFlow
    tf.keras.utils.get_custom_objects()['CustomNonPaddingTokenLoss'] = CustomNonPaddingTokenLoss
    loaded_model = Sequential()
    loaded_model.add(TFSMLayer("./resources/ner_model", call_endpoint='serving_default'))


    all_predicted_tag_ids = []
    
    for x, _ in data:
        print("Input Tensor Info:")
        print("Data Type:", x.dtype)
        print("Shape:", x.shape)
        x = tf.cast(x, tf.int64)
        output = loaded_model(x, training=False)
        predictions = np.argmax(output, axis=-1)
        predictions = np.reshape(predictions, [-1])
        all_predicted_tag_ids.append(predictions)

    all_predicted_tag_ids = np.concatenate(all_predicted_tag_ids)

    ner_labels = ["[PAD]", "N", "M", "other"]
    mapping =  dict(zip(range(len(ner_labels)), ner_labels))
    predicted_tags = [mapping[tag] for tag in all_predicted_tag_ids]

    return predicted_tags

sample_input = "Hello world, my name is John, I live in New York, my birthday is 10/02/1990."

result = prediction(format_datatype(sample_input))
print(result)

• 固定padded_batch的输出类型：padded_batch默认会将张量转为float32，需显式指定padding值的类型为int64，避免隐式转换：
  修改format_datatype中的padded_batch调用：

  .padded_batch(1, padding_values=(tf.constant(0, dtype=tf.int64), tf.constant(0, dtype=tf.int64)))
  

• 在StringLookup阶段指定类型：创建lookup_layer时直接指定输出类型为int64，从源头固定张量类型：

  lookup_layer = keras.layers.StringLookup(vocabulary=loaded_list, dtype=tf.int64)
  

• 检查模型导出的签名：确认Colab导出模型时的输入签名类型为int64，若签名有误需重新导出。可通过以下代码查看签名：

  import tensorflow as tf
  model = tf.saved_model.load("./resources/ner_model")
  print(model.signatures['serving_default'].inputs)
  

• 简化预处理流程：跳过tf.data.Dataset的复杂处理，直接生成int64类型的输入张量：

  def format_datatype(data):
      tokens = [re.sub(r'[;,]', '', d) for d in data.split(' ')]
      with open('./resources/vocabulary.pkl', 'rb') as f:
          loaded_list = pickle.load(f)
      lookup_layer = keras.layers.StringLookup(vocabulary=loaded_list, dtype=tf.int64)
      token_ids = lookup_layer(tokens)
      # 增加batch维度并padding（max_seq_len需替换为模型要求的序列长度）
      token_ids = tf.expand_dims(token_ids, 0)
      token_ids = tf.pad(token_ids, [[0,0], [0, 32 - len(token_ids)]], constant_values=0)
      return token_ids
  

内容的提问来源于stack exchange，提问作者Yin Jie