FlaxLayer

[源文件]

FlaxLayer

keras.layers.FlaxLayer(module, method=None, variables=None, **kwargs)

封装 Flax 模块的 Keras 层。

当使用 JAX 作为 Keras 的后端时,该层可以在 Keras 中以 flax.linen.Module 实例的形式使用 Flax 组件。

用于前向传播的模块方法可以通过 method 参数指定,默认情况下为 __call__。此方法必须接受以下名称完全一致的参数:

  • 如果方法绑定到模块(对于默认的 __call__ 而言就是这种情况),则为 self;否则为 module,用于传递模块。
  • inputs:模型的输入,一个 JAX 数组或由数组组成的 PyTree
  • training (可选):一个指定我们是否处于训练模式或推理模式的参数,在训练模式下会传递 True

FlaxLayer 会自动处理模型的非训练状态和所需的 RNG。请注意,flax.linen.Module.apply()mutable 参数被设置为 DenyList(["params"]),因此假定所有不在 "params" 集合中的变量都是非训练权重。

此示例展示了如何从具有默认 __call__ 方法且无训练参数的 Flax Module 创建一个 FlaxLayer

class MyFlaxModule(flax.linen.Module):
    @flax.linen.compact
    def __call__(self, inputs):
        x = inputs
        x = flax.linen.Conv(features=32, kernel_size=(3, 3))(x)
        x = flax.linen.relu(x)
        x = flax.linen.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = x.reshape((x.shape[0], -1))  # flatten
        x = flax.linen.Dense(features=200)(x)
        x = flax.linen.relu(x)
        x = flax.linen.Dense(features=10)(x)
        x = flax.linen.softmax(x)
        return x

flax_module = MyFlaxModule()
keras_layer = FlaxLayer(flax_module)

此示例展示了如何封装模块方法以符合所需的签名。这允许具有多个输入参数以及一个名称和值不同的训练参数。此外,此示例还展示了如何使用一个未绑定到模块的函数。

class MyFlaxModule(flax.linen.Module):
    @flax.linen.compact
    def forward(self, input1, input2, deterministic):
        ...
        return outputs

def my_flax_module_wrapper(module, inputs, training):
    input1, input2 = inputs
    return module.forward(input1, input2, not training)

flax_module = MyFlaxModule()
keras_layer = FlaxLayer(
    module=flax_module,
    method=my_flax_module_wrapper,
)

参数

  • moduleflax.linen.Module 或其子类的实例。
  • method:要调用的模型方法。这通常是 Module 中的一个方法。如果未提供,则使用 __call__ 方法。method 也可以是一个未在 Module 中定义的函数,在这种情况下,它必须将 Module 作为第一个参数。它同时用于 Module.initModule.apply。详细信息请参阅 flax.linen.Module.apply()method 参数文档。
  • variables:一个 dict,包含模块的所有变量,格式与 flax.linen.Module.init() 返回的格式相同。它应该包含一个 "params" 键,并在适用时包含其他键用于非训练状态的变量集合。这允许传递已训练的参数和学习到的非训练状态,或控制初始化。如果传递 None,则在构建时调用模块的 init 函数来初始化模型的变量。