Flax深度学习库构建简单的神经网络模型

Flax 是一个构建在 JAX（Just Another eXperimental library）之上的深度学习库，它旨在提供高级的、模块化的接口，以便更轻松地定义、训练和部署神经网络模型。 Flax 和 Haiku 是两个相互关联的深度学习库，它们都是基于 JAX 构建的，旨在提供高级的模块化接口。 Utilities可以在Haiku和Flax之间无缝移动。

from flax import linen as flax_nn
import jax
import jax.numpy as jnp

class SimpleModel(flax_nn.Module):
    hidden_size: int
    output_size: int
    def setup(self):
        # 在 setup 方法中定义模块的初始化逻辑
        self.hidden = nn.Dense(features=self.hidden_size)
        self.output = nn.Dense(features=self.output_size)
        
    def __call__(self, x):
        # 前向传播逻辑，包括全连接和激活函数
        x = self.hidden(x)
        x= jax.nn.relu(x)
        x = self.output(x)
        return jax.nn.sigmoid(x)

# 创建模型实例,继承自flax_nn.Module，不需要创建函数
# 如果是haiku模型（继承自hk.Module），需要创建函数，再对函数进行transform转化（hk.transform）
model = SimpleModel(hidden_size=64, output_size=10)
#print(model)

# 初始化模型参数
rng = jax.random.PRNGKey(42)
input_data = jax.random.normal(rng, (1, 128))
params = model.init(rng, input_data)

# 执行前向传播，二个参数（haiku transform构建的模块，其apply方法需三个参数）
output = model.apply(params, input_data)

print("Input shape:", input_data.shape)
print("Output shape:", output.shape)

参考：

https://github.com/google/flax

https://dm-haiku.readthedocs.io/en/latest/notebooks/flax.html