如何在json模型导出中保留自定义层名称并避免“TensorFlowPlayer”

Question

我正在训练一个序列tf.keras模型，我想将其转换为tfjs格式，由描述层的model.json文件和二进制权重文件组成，以便将其部署到网站上进行推断

我的模型中有两个层是自定义层，因为在tf.keras.layers中没有预定义的合适层来执行此工作。这是我的模型代码的模拟版本：

import tensorflow as tf

class SliceChannelsLayer(tf.keras.layers.Layer):
    ... (model implementation) ...

class L1NormLayer(tf.keras.layers.Layer):
    ... (model implementation) ...

def sequential_model():
    inputs = tf.keras.Input(shape=(30, 36, 36, 6), batch_size=None)
    outputs_a = SliceChannelsLayer(start=0, end=3)(inputs)
    outputs_b = SliceChannelsLayer(start=3, end=6)(inputs)
    ...
    other Keras layers
    ...
    attention = L1NormLayer(1)(attention)
    outputs_motion = tf.keras.layers.Multiply()([outputs_a, attention])
    return tf.keras.Model(inputs, outputs_b)

model = sequential_model()

在我的JavaScript代码中，我将两个自定义层实现为outlined in the example given in tfjs-examples：

class SliceChannelsLayer extends tf.layers.Layer {
  ... (model implementation) ...
}

class L1NormLayer extends tf.layers.Layer {
  ... (model implementation) ...
}

运行tfjs.converters.save_keras_model(model, "tfjs_export")时，两个自定义层在model.json:TensorFlowOpLayer中接收相同的名称

但是，我需要将JavaScript层实现的名称与这些名称相匹配，这在多个层接收相同名称时是不可能的。 当我手动编辑model.json文件以TensorFlowOpLayer替换为SliceChannelsLayer或L1NormLayer时，我的代码可以工作

我的问题：如何避免在导出后必须手动更改model.json中的图层名称

编辑：可使用以下代码复制该行为：

import tensorflow as tf
import tensorflowjs as tfjs

class L1NormLayer(tf.keras.layers.Layer):
    """L1NormLayer"""
    def __init__(self, axis, **kwargs):
        super(L1NormLayer, self).__init__()
        self.axis = axis
    def __call__(self, inputs):
        inputs, _ = tf.linalg.normalize(inputs, ord=1, axis=self.axis)
        return inputs

class SliceChannelsLayer(tf.keras.layers.Layer):
    """SliceChannelsLayer"""
    def __init__(self, start, end, **kwargs):
        super(SliceChannelsLayer, self).__init__()
        self.start = start
        self.end = end
    def __call__(self, inputs):
        inputs = inputs[:, :, :, :, self.start:self.end]
        return inputs

inputs = tf.keras.Input(shape=(29, 36, 36, 6), batch_size=1)
outputs_a = SliceChannelsLayer(start=0, end=3, name="SliceChannels")(inputs)
outputs_b = SliceChannelsLayer(start=3, end=6, name="SliceChannels")(inputs)
attention = tf.keras.layers.TimeDistributed(
    tf.keras.layers.Conv2D(
        filters=1, kernel_size=(1, 1), activation="sigmoid"))(
            outputs_b)
attention = L1NormLayer(1, name="L1Norm")(attention)
outputs_a = tf.keras.layers.Multiply()([outputs_b, attention])
outputs_a = tf.keras.layers.Dense(units=1, activation="linear")(outputs_a)
model = tf.keras.Model(inputs=inputs, outputs=outputs_a)

model.compile(loss='mean_squared_error',
              optimizer=tf.keras.optimizers.RMSprop())

model.summary()

tfjs.converters.save_keras_model(model, "tfjs")

Answer 1

__call__方法应改为call