`haiku.get_state` 和 `haiku.set_state` 是 DeepMind 的 Haiku 库中的两个核心函数,用于在 JAX 框架中管理可变状态。这些函数允许你在 JAX 的纯函数式编程环境中使用状态(如批归一化层的均值和方差、RNN 的隐藏状态等),并在前向传播过程中读取和更新这些状态。
`haiku.get_state`
- 功能: `haiku.get_state` 用于获取模型的当前状态。状态可以是任何需要在模型运行期间更新的值。
- 常见用法: 当你需要在某个位置保存一个可变值(如均值、方差等)并在每次前向传播时读取它时,可以使用 `get_state`。
语法
state = hk.get_state(name, shape, dtype, init)
- `name`: 状态的名称,字符串类型。
- `shape`: 状态的形状,通常为一个元组。
- `dtype`: 状态的数据类型,如 `jnp.float32`。
- `init`: 初始化函数,用于在状态不存在时创建初始值。
示例
def my_module(x):
# 初始化一个状态变量
state = hk.get_state("my_state", shape=[], dtype=jnp.float32, init=jnp.zeros)
# 使用这个状态
y = x + state
return y
`haiku.set_state`
- 功能: `haiku.set_state` 用于更新模型的状态。你可以在前向传播过程中使用这个函数来更新状态的值。
- 常见用法: 在模型的某个操作后,将新的值存储回状态中,供下一次前向传播使用。
语法
hk.set_state(name, value)
- `name`: 状态的名称,字符串类型。
- `value`: 要设置的新状态值,应该与 `get_state` 获取的状态形状和类型匹配。
示例
def my_module(x):
# 获取状态
state = hk.get_state("my_state", shape=[], dtype=jnp.float32, init=jnp.zeros)
# 更新状态
new_state = state + 1
hk.set_state("my_state", new_state)
y = x + state
return y
完整示例
以下是一个完整的例子,演示如何使用 `get_state` 和 `set_state` 来管理和更新模型的状态:
import haiku as hk
import jax.numpy as jnp
import jax
# 定义一个模块,使用 Haiku 的 get_state 和 set_state
def my_module(x):
# 获取状态,初始化为0
state = hk.get_state("counter", shape=[], dtype=jnp.float32, init=jnp.zeros)
# 更新状态,每次调用时将状态增加1
new_state = state + 1
hk.set_state("counter", new_state)
# 返回输入加上当前状态
return x + state
# 转换为 Haiku 模块
def forward_fn(x):
return my_module(x)
forward = hk.transform_with_state(forward_fn)
# 初始化参数和状态
params, state = forward.init(jax.random.PRNGKey(42), jnp.array(1.0))
# 应用模型
for _ in range(5):
output, state = forward.apply(params, state, None, jnp.array(1.0))
print(f"Output: {output}, State: {state['my_module']['counter']}")
输出示例
Output: 1.0, State: 1.0
Output: 2.0, State: 2.0
Output: 3.0, State: 3.0
Output: 4.0, State: 4.0
Output: 5.0, State: 5.0
在这个示例中,每次调用模型时,状态 `counter` 都会增加1,反映出在前向传播过程中状态的变化。
`hk.get_parameter`
- 功能: `hk.get_parameter` 用于获取一个模型参数。参数通常是可训练的变量,如神经网络中的权重或偏置。该函数允许你定义或获取已存在的参数,并在模型的前向传播中使用这些参数。
语法
param = hk.get_parameter(name, shape, dtype, init)
- `name`: 参数的名称,字符串类型。
- `shape`: 参数的形状,通常为一个元组。
- `dtype`: 参数的数据类型,如 `jnp.float32`。
- `init`: 初始化函数,用于在参数不存在时创建初始值。
示例
以下是一个简单的例子,演示如何使用 `hk.get_parameter` 在 Haiku 模块中创建和获取模型参数:
import haiku as hk
import jax.numpy as jnp
import jax
def my_module(x):
# 创建或获取一个名为 'w' 的参数
w = hk.get_parameter("w", shape=[x.shape[-1], 1], dtype=jnp.float32, init=hk.initializers.RandomNormal())
# 应用参数到输入数据上
y = jnp.dot(x, w)
return y
# 转换为 Haiku 模块
def forward_fn(x):
return my_module(x)
forward = hk.transform(forward_fn)
# 初始化参数
params = forward.init(jax.random.PRNGKey(42), jnp.array([[1.0, 2.0, 3.0]]))
# 应用模型
output = forward.apply(params, None, jnp.array([[1.0, 2.0, 3.0]]))
print(output)
输出示例
[[-2.2916665]] # 随机初始化的权重乘以输入向量的结果
说明
- 创建参数: `hk.get_parameter` 在参数第一次被请求时,会使用 `init` 函数来初始化参数。
- 获取参数: 如果参数已经存在(即在同一层的后续调用中),`hk.get_parameter` 会直接返回该参数,而不会重新初始化。