当自定义对象作为 Tensorflow 中的梯度反转层时加载模型中的内容（域适应）

Question

所以，这是域适应模型的示例代码，我只想保存模型并加载它，

@tf.custom_gradient
def grad_reverse(x):
    y = tf.identity(x)
    def custom_grad(dy):
        return -dy
    return y, custom_grad

class GradReverse(tf.keras.layers.Layer):
    def __init__(self):
        super().__init__(name="grl")

    def call(self, x):
        return grad_reverse(x)


def get_adaptable_network(input_shape=x_source_train.shape[1:]):
    
    inputs = Input(shape=input_shape)
    x = Conv2D(32, 5, padding='same', activation='relu', name='conv2d_1')(inputs)
    x = MaxPool2D(pool_size=2, strides=2, name='max_pooling2d_1')(x)
    x = Conv2D(48, 5, padding='same', activation='relu', name='conv2d_2')(x)
    x = MaxPool2D(pool_size=2, strides=2, name='max_pooling2d_2')(x)
    features = Flatten(name='flatten_1')(x)
    x = Dense(100, activation='relu', name='dense_digits_1')(features)
    x = Dense(100, activation='relu', name='dense_digits_2')(x)
    digits_classifier = Dense(10, activation="softmax", name="digits_classifier")(x)

    domain_branch = Dense(100, activation="relu", name="dense_domain")(GradReverse()(features))
    domain_classifier = Dense(1, activation="sigmoid", name="domain_classifier")(domain_branch)

    return Model(inputs=inputs, outputs=[digits_classifier, domain_classifier])

model = get_adaptable_network()
model.summary()

# download the model in computer for later use
model.save('DA_MNIST_to_MNIST_m.h5')

from tensorflow import keras
model = keras.models.load_model('DA_MNIST_to_MNIST_m.h5',custom_objects={'?':? })

我不确定在 custom_objects 部分放什么，因为在 tensorflow 中为域自适应实现了自定义梯度反转层。当我加载模型时，出现错误：

/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/utils/generic_utils.py in class_and_config_for_serialized_keras_object(config, module_objects, custom_objects, printable_module_name)
    294   cls = get_registered_object(class_name, custom_objects, module_objects)
    295   if cls is None:
--> 296     raise ValueError('Unknown ' + printable_module_name + ': ' + class_name)
    297 
    298   cls_config = config['config']

ValueError: Unknown layer: GradReverse

我正在对 MNIST_M 域进行 MNIST 适配，任何帮助都会很有用！

Answer 1

我想通了，我需要用 **kwargs 更改 GradReverse 层的初始化函数，然后这个对象将接受我没有包含的任何其他关键字参数。

class GradReverse(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(name="grl")

    def call(self, x):
        return grad_reverse(x)

在负载模型中，我们可以使用这个，

from tensorflow import keras
model = keras.models.load_model('DA_MNIST_to_MNIST_m.h5',custom_objects={'GradReverse':GradReverse})

当自定义对象作为 Tensorflow 中的梯度反转层时加载模型中的内容（域适应）

What to put in load model when there is custom objects as gradient reversal layer in Tensorflow (Domain Adaptation )

python-3.x

tensorflow

keras

deep-learning

transfer-learning