迁移到普通字典

在调用 .init、.init_with_output 和 .apply Module 方法时，Flax 将从返回 FrozenDicts 迁移到返回普通字典。

原始问题在此处概述。

本指南展示了一些常见的升级模式。

实用函数

FrozenDicts 是不可变的字典，它实现了额外的 4 个方法

• copy

• pop

• pretty_repr

• unfreeze

为了适应普通字典的更改，请将 FrozenDict 方法的使用替换为 flax.core.frozen_dict 中等效的实用函数。这些实用函数模仿其对应的 FrozenDict 方法的行为，并且可以在 FrozenDicts 或普通字典上调用。以下是实用函数和示例升级模式

copy

仅 FrozenDict

variables = variables.copy(add_or_replace={'other_variables': other_variables})

既有 FrozenDict 也有普通字典

variables = flax.core.copy(variables, add_or_replace={'other_variables': other_variables})

pop

仅 FrozenDict

state, params = variables.pop('params')

既有 FrozenDict 也有普通字典

state, params = flax.core.pop(variables, 'params')

pretty_repr

仅 FrozenDict

str_repr = variables.pretty_repr()

既有 FrozenDict 也有普通字典

str_repr = flax.core.pretty_repr(variables)

unfreeze

仅 FrozenDict

variables = variables.unfreeze()

既有 FrozenDict 也有普通字典

variables = flax.core.unfreeze(variables)

修改配置值

设置了一个临时特性标志 flax_return_frozendict 以帮助进行迁移。要在运行时切换返回 FrozenDict 和普通字典变量的行为，请在代码中运行 flax.config.update('flax_return_frozendict', <BOOLEAN_VALUE>)。

例如

x = jnp.empty((1,3))

flax.config.update('flax_return_frozendict', True) # set Flax to return FrozenDicts
variables = nn.Dense(5).init(jax.random.key(0), x)

assert isinstance(variables, flax.core.FrozenDict)

flax.config.update('flax_return_frozendict', False) # set Flax to return regular dicts
variables = nn.Dense(5).init(jax.random.key(0), x)

assert isinstance(variables, dict)

或者，可以直接在 Flax 源代码中修改环境变量 flax_return_frozendict（在此处找到）。

迁移状态

截至 2023 年 7 月 19 日，flax_return_frozendict 设置为 False（请参阅 #3193），这意味着 Flax 将默认从 0.7.1 版本开始返回普通字典。可以临时将此标志翻转为 True，让 Flax 返回 Frozendicts。但是，此特性标志最终将在未来删除。