使用自定义张量作为变量的 TensorFlow 2.0 Keras 层

Question

在 TF 1.x 中，可以使用自定义变量构建层。下面是一个例子：

import numpy as np
import tensorflow as tf

def make_custom_getter(custom_variables):
    def custom_getter(getter, name, **kwargs):
        if name in custom_variables:
            variable = custom_variables[name]
        else:
            variable = getter(name, **kwargs)
        return variable
    return custom_getter

# Make a custom getter for the dense layer variables.
# Note: custom variables can result from arbitrary computation;
#       for the sake of this example, we make them just constant tensors.
custom_variables = {
    "model/dense/kernel": tf.constant(
        np.random.rand(784, 64), name="custom_kernel", dtype=tf.float32),
    "model/dense/bias": tf.constant(
        np.random.rand(64), name="custom_bias", dtype=tf.float32),
}
custom_getter = make_custom_getter(custom_variables)

# Compute hiddens using a dense layer with custom variables.
x = tf.random.normal(shape=(1, 784), name="inputs")
with tf.variable_scope("model", custom_getter=custom_getter):
    Layer = tf.layers.Dense(64)
    hiddens = Layer(x)

print(Layer.variables)

构造的密集层的打印变量将是我们在custom_variables字典中指定的自定义张量：

[<tf.Tensor 'custom_kernel:0' shape=(784, 64) dtype=float32>, <tf.Tensor 'custom_bias:0' shape=(64,) dtype=float32>]

这允许我们创建custom_variables直接使用提供的张量作为权重的层/模型，以便我们可以进一步区分层/模型的输出与custom_variables可能依赖的任何张量（对于实现调制子函数的功能特别有用）nets、参数生成、元学习等）。

变量范围用于使用自定义 getter 轻松嵌套所有图形构建内部范围，并在提供的张量之上构建模型作为它们的参数。由于 TF 2.0 中不再建议使用会话和变量范围（并且所有这些低级内容都移至tf.compat.v1），因此使用 Keras 和 TF 2.0 实现上述内容的最佳实践是什么？

（GitHub 上的相关问题。）

Answer 1

下面是一个通用解决方案，适用于 TF2 中的任意 Keras 模型。

首先，我们需要定义一个辅助函数canonical_variable_name和一个custom_make_variable具有以下签名的上下文管理器（请参阅元块库中的实现）。

def canonical_variable_name(variable_name: str, outer_scope: str):
    """Returns the canonical variable name: `outer_scope/.../name`."""
    # ...

@contextlib.contextmanager
def custom_make_variable(
    canonical_custom_variables: Dict[str, tf.Tensor], outer_scope: str
):
    """A context manager that overrides `make_variable` with a custom function.

    When building layers, Keras uses `make_variable` function to create weights
    (kernels and biases for each layer). This function wraps `make_variable` with
    a closure that infers the canonical name of the variable being created (of the
    form `outer_scope/.../var_name`) and looks it up in the `custom_variables` dict
    that maps canonical names to tensors. The function adheres the following logic:

    * If there is a match, it does a few checks (shape, dtype, etc.) and returns
      the found tensor instead of creating a new variable.
    * If there is a match but checks fail, it throws an exception.
    * If there are no matching `custom_variables`, it calls the original
      `make_variable` utility function and returns a newly created variable.
    """
    # ...

使用这些函数，我们可以使用自定义张量作为变量来创建任意 Keras 模型：

import numpy as np
import tensorflow as tf

canonical_custom_variables = {
    "model/dense/kernel": tf.constant(
        np.random.rand(784, 64), name="custom_kernel", dtype=tf.float32),
    "model/dense/bias": tf.constant(
        np.random.rand(64), name="custom_bias", dtype=tf.float32),
}

# Compute hiddens using a dense layer with custom variables.
x = tf.random.normal(shape=(1, 784), name="inputs")
with custom_make_variable(canonical_custom_variables, outer_scope="model"):
    Layer = tf.layers.Dense(64)
    hiddens = Layer(x)

print(Layer.variables)