Haiku 库中的get_state和set_state及get_parameter函数

`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` 会直接返回该参数，而不会重新初始化。