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，不需要创建函数
# 如果