如何保留自定义层名称并避免 json 模型导出中的 "TensorFlowOpLayer"

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。

如何保留自定义层名称并避免 json 模型导出中的 "TensorFlowOpLayer"

How to retain custom layer names and avoid "TensorFlowOpLayer" in json model export

javascript

python

tensorflow

tensorflow.js