Flax模型实例化开销：重复创建网络是否成本高昂？

你的示例代码：

class QNetwork(nn.Module):
    dim: int

    @nn.compact
    def __call__(self, x):
        x = nn.Dense(120)(x)
        x = nn.relu(x)
        x = nn.Dense(84)(x)
        x = nn.relu(x)
        x = nn.Dense(self.dim)(x)
        return x

q = QNetwork(4)

# do stuff

q = QNetwork(4)

# do stuff

q = QNetwork(4)

# do stuff

你的理解是对的，重复创建Flax网络实例的开销确实很低，原因如下：

• Flax的nn.Module本质是无状态的架构定义容器，实例化时只保存网络的结构参数（比如这里的dim=4、各层的神经元数量配置），并不存储模型的权重参数。权重是在首次调用模块传入数据，或是调用init方法时才会生成，且这些权重是独立于模块实例的JAX数组。
• 重复执行QNetwork(4)只是创建一个轻量的配置对象，不会涉及内存密集的权重分配或复制操作，因此开销几乎可以忽略。

需要注意的是：如果每次创建实例后都重新调用init生成新权重，那这部分会产生权重初始化的开销，但这是权重生成的成本，和模块实例本身的创建无关。如果只是重复创建实例但复用已有的权重，完全不会有额外的高开销。

内容的提问来源于stack exchange，提问作者joel