JaxLayer
类keras.layers.JaxLayer(
call_fn, init_fn=None, params=None, state=None, seed=None, **kwargs
)
包装 JAX 模型的 Keras 层。
当使用 JAX 作为 Keras 的后端时,此层可实现在 Keras 中使用 JAX 组件。
此层接受函数形式的 JAX 模型,即 call_fn
,该函数必须包含以下参数且名称必须完全一致:
params
:模型的可训练参数。state
(可选):模型的不可训练状态。如果模型没有不可训练状态,则可以省略。rng
(可选):一个 jax.random.PRNGKey
实例。如果模型在训练或推断期间不需要 RNG,则可以省略。inputs
:模型的输入,一个 JAX 数组或数组的 PyTree
。training
(可选):一个指定是否处于训练模式或推断模式的参数,在训练模式下传入 True
。如果模型在训练模式和推断模式下行为相同,则可以省略。inputs
参数是必需的。模型的输入必须通过一个参数提供。如果 JAX 模型将多个输入作为单独的参数,则必须将它们组合成一个单一的结构,例如在 tuple
或 dict
中。
模型 params
和 state
的初始化可以由此层处理,在这种情况下,必须提供 init_fn
参数。这允许模型以正确的形状动态初始化。另外,如果形状已知,可以使用 params
参数和可选的 state
参数来创建已经初始化的模型。
如果提供 init_fn
函数,它必须包含以下参数且名称必须完全一致:
rng
:一个 jax.random.PRNGKey
实例。inputs
:一个 JAX 数组或具有占位符值的数组的 PyTree
,以提供输入形状。training
(可选):一个指定是否处于训练模式或推断模式的参数。True
始终传递给 init_fn
。无论 call_fn
是否有 training
参数,都可以省略。对于具有不可训练状态的 JAX 模型:
call_fn
必须有一个 state
参数call_fn
必须返回一个包含模型输出和模型新的不可训练状态的 tuple
init_fn
必须返回一个包含模型初始可训练参数和初始不可训练状态的 tuple
。此代码展示了具有不可训练状态的模型 call_fn
和 init_fn
签名的可能组合。在此示例中,模型在 call_fn
中有一个 training
参数和一个 rng
参数。
def stateful_call(params, state, rng, inputs, training):
outputs = ...
new_state = ...
return outputs, new_state
def stateful_init(rng, inputs):
initial_params = ...
initial_state = ...
return initial_params, initial_state
对于没有不可训练状态的 JAX 模型:
call_fn
不得有 state
参数call_fn
必须只返回模型的输出init_fn
必须只返回模型的初始可训练参数。此代码展示了没有不可训练状态的模型 call_fn
和 init_fn
签名的可能组合。在此示例中,模型在 call_fn
中没有 training
参数,也没有 rng
参数。
def stateless_call(params, inputs):
outputs = ...
return outputs
def stateless_init(rng, inputs):
initial_params = ...
return initial_params
如果模型的签名与 JaxLayer
所需的签名不同,可以轻松编写一个包装方法来调整参数。此示例展示了一个模型,它有多个输入作为单独的参数,期望 dict
中有多个 RNG,并且有一个 deterministic
参数,其含义与 training
相反。为了符合要求,输入通过 tuple
组合成一个单一结构,RNG 被拆分并用于填充预期的 dict
,布尔标志被否定。
def my_model_fn(params, rngs, input1, input2, deterministic):
...
if not deterministic:
dropout_rng = rngs["dropout"]
keep = jax.random.bernoulli(dropout_rng, dropout_rate, x.shape)
x = jax.numpy.where(keep, x / dropout_rate, 0)
...
...
return outputs
def my_model_wrapper_fn(params, rng, inputs, training):
input1, input2 = inputs
rng1, rng2 = jax.random.split(rng)
rngs = {"dropout": rng1, "preprocessing": rng2}
deterministic = not training
return my_model_fn(params, rngs, input1, input2, deterministic)
keras_layer = JaxLayer(my_model_wrapper_fn, params=initial_params)
JaxLayer
允许使用 Haiku 组件,形式为 haiku.Module
。这是通过按照 Haiku 模式转换模块来实现的,然后将 module.apply
传递给 call_fn
参数,如果需要,将 module.init
传递给 init_fn
参数。
如果模型有不可训练状态,则应使用 haiku.transform_with_state
进行转换。如果模型没有不可训练状态,则应使用 haiku.transform
进行转换。此外,如果模块在“apply”中不使用 RNG,则可以选择使用 haiku.without_apply_rng
进行转换。
以下示例展示了如何从使用 hk.next_rng_key()
通过随机数生成器并接受训练位置参数的 Haiku 模块创建 JaxLayer
:
class MyHaikuModule(hk.Module):
def __call__(self, x, training):
x = hk.Conv2D(32, (3, 3))(x)
x = jax.nn.relu(x)
x = hk.AvgPool((1, 2, 2, 1), (1, 2, 2, 1), "VALID")(x)
x = hk.Flatten()(x)
x = hk.Linear(200)(x)
if training:
x = hk.dropout(rng=hk.next_rng_key(), rate=0.3, x=x)
x = jax.nn.relu(x)
x = hk.Linear(10)(x)
x = jax.nn.softmax(x)
return x
def my_haiku_module_fn(inputs, training):
module = MyHaikuModule()
return module(inputs, training)
transformed_module = hk.transform(my_haiku_module_fn)
keras_layer = JaxLayer(
call_fn=transformed_module.apply,
init_fn=transformed_module.init,
)
参数
None
,则必须提供 params
和/或 state
。PyTree
。这允许传递训练过的参数或控制初始化。如果 params
和 state
都为 None
,则在构建时调用 init_fn
以初始化模型的可训练参数。PyTree
。这允许传递学习到的状态或控制初始化。如果 params
和 state
都为 None
,并且 call_fn
接受 state
参数,则在构建时调用 init_fn
以初始化模型的不可训练状态。keras.DTypePolicy
。可选。默认为默认策略。