交互式在线版本：

Haiku 中的参数共享#

简介#

在 Haiku 中，参数重用由模块实例名称唯一确定，即，如果一个模块实例与另一个模块实例具有相同的名称，则它们共享参数。

除非另有说明，模块名称由 Haiku 根据模块类名称自动确定（遵循 TensorFlow 1 和 Sonnet V1 中建立的模式）。更详细地说，模块命名遵循以下规则

模块名称在模块实例构造时分配。除非模块实例名称作为参数提供给构造函数，否则 Haiku 从当前的模块类名称生成一个名称（基本上：to_snake_case(CurrentClassName)）。
如果模块实例名称不以 _N 结尾（其中 N 是一个数字），并且另一个具有相同名称的模块实例已经存在，则 Haiku 会在新模块实例名称的末尾添加一个递增的数字（例如 module_1）。
当两个模块嵌套时（即，在一个模块的类定义内部构造一个模块实例），则内部模块名称将以外部模块名称和（可能，见下一点）外部模块当前调用的方法为前缀。构造函数（即 __init__）被波浪号 ~ 符号替换。
如果调用方法名称是 __call__，这将忽略（方法名称将仅以外部模块名称为前缀）。
当存在多层嵌套时，先前的规则应用于每个嵌套级别，并且每个内部模块名称都基于模块名称和调用方法名称，该模块紧接在调用层次结构中的当前模块之前。

让我们看看这如何通过一个实际示例来工作。

扁平模块（无嵌套）#

本节介绍模块未嵌套时的参数共享。

[4]:

#@title Imports and accessory functions
import functools
import haiku as hk
import jax
import jax.numpy as jnp


def parameter_shapes(params):
  """Make printing parameters a little more readable."""
  return jax.tree_util.tree_map(lambda p: p.shape, params)


def transform_and_print_shapes(fn, x_shape=(2, 3)):
  """Print name and shape of the parameters."""
  rng = jax.random.PRNGKey(42)
  x = jnp.ones(x_shape)

  transformed_fn = hk.transform(fn)
  params = transformed_fn.init(rng, x)
  print('\nThe name and shape of the parameters are:')
  print(parameter_shapes(params))

def assert_all_equal(params_1, params_2):
  assert all(jax.tree_util.tree_leaves(
      jax.tree_util.tree_map(lambda a, b: (a == b).all(), params_1, params_2)))

[6]:

w_init = hk.initializers.TruncatedNormal(stddev=1)

class SimpleModule(hk.Module):
  """A simple module class with one variable."""

  def __init__(self, output_channels, name=None):
    super().__init__(name)
    assert isinstance(output_channels, int)
    self._output_channels = output_channels

  def __call__(self, x):
    w_shape = (x.shape[-1], self._output_channels)
    w = hk.get_parameter("w", w_shape, x.dtype, init=w_init)
    return jnp.dot(x, w)

[ ]:

def f(x):
  # This instance will be named `a_simple_module`.
  simple = SimpleModule(output_channels=2)
  simple_out = simple(x)  # implicitly calls module_install.__call__()
  print(f'The name assigned to "simple" is: "{simple.module_name}".')
  return simple_out

transform_and_print_shapes(f)

WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

The name assigned to "simple" is: "simple_module".

The name and shape of the parameters are:
{'simple_module': {'w': (3, 2)}}

太棒了！在这里我们看到，如果我们创建一个 SimpleModule 实例并且不指定名称，Haiku 会为其分配名称 a_simple_module。这也反映在与模块关联的参数中。

如果我们实例化两次 SimpleModule 会发生什么？Haiku 是否为两个实例分配相同的名称？

[ ]:

def f(x):
  # This instance will be named `a_simple_module`.
  simple_one = SimpleModule(output_channels=2)
  # This instance will be named `a_simple_module_1`.
  simple_two = SimpleModule(output_channels=2)
  first_out = simple_one(x)
  second_out = simple_two(x)
  print(f'The name assigned to "simple_one" is: "{simple_one.module_name}".')
  print(f'The name assigned to "simple_two" is: "{simple_two.module_name}".')
  return first_out + second_out

transform_and_print_shapes(f)

The name assigned to "simple_one" is: "simple_module".
The name assigned to "simple_two" is: "simple_module_1".

The name and shape of the parameters are:
{'simple_module': {'w': (3, 2)}, 'simple_module_1': {'w': (3, 2)}}

正如预期的那样，Haiku 足够智能，可以区分两个实例并避免意外的参数共享：第二个实例被命名为 a_simple_module_1，并且每个实例都有自己的一组参数。很好！

但是，如果我们想要共享参数呢？在这种情况下，我们必须仅实例化模块一次并调用它多次。让我们看看这是如何工作的

[ ]:

def f(x):
  # This instance will be named `a_simple_module`.
  simple_one = SimpleModule(output_channels=2)
  first_out = simple_one(x)
  second_out = simple_one(x)  # share parameters w/ previous call
  print(f'The name assigned to "simple_one" is: "{simple_one.module_name}".')
  return first_out + second_out

transform_and_print_shapes(f)

The name assigned to "simple_one" is: "simple_module".

The name and shape of the parameters are:
{'simple_module': {'w': (3, 2)}}

嵌套模块#

在本节中，我们将看到当我们将一个 hk.Module 嵌套到另一个中时会发生什么。

[ ]:

class NestedModule(hk.Module):
  """A module class with a nested module created in the constructor."""

  def __init__(self, output_channels, name=None):
    super().__init__(name)
    assert isinstance(output_channels, int)
    self._output_channels = output_channels
    self.inner_simple = SimpleModule(self._output_channels)

  def __call__(self, x):
    w_shape = (x.shape[-1], self._output_channels)
    # Another variable that is also called `w`.
    w = hk.get_parameter("w", w_shape, x.dtype, init=w_init)
    return jnp.dot(x, w) + self.inner_simple(x)

[ ]:

def f(x):
  # This will be named `a_nested_module` and the SimpleModule instance created
  # inside it will be named `a_nested_module/a_simple_module`.
  nested = NestedModule(output_channels=2)
  nested_out = nested(x)
  print('The name assigned to outer module (i.e., "nested") is: '
        f'"{nested.module_name}".')
  print('The name assigned to the inner module (i.e., inside "nested") is: "'
        f'{nested.inner_simple.module_name}".')
  return nested_out

transform_and_print_shapes(f)

The name assigned to outer module (i.e., "nested") is: "nested_module".
The name assigned to the inner module (i.e., inside "nested") is: "nested_module/~/simple_module".

The name and shape of the parameters are:
{'nested_module': {'w': (3, 2)}, 'nested_module/~/simple_module': {'w': (3, 2)}}

正如预期的那样，内部模块名称取决于：（a）外部模块名称；以及（b）正在调用的外部模块的方法。

另请注意，外部模块的构造函数名称 __init__ 如何在参数名称中被 ~ 替换。如果内部模块实例是在外部模块的 __call__ 方法内部创建的，则内部模块实例名称将为 'a_nested_module/a_simple_module'。

在本例中，我们从头开始定义了所有模块，但对于 Haiku 中定义的任何模块和网络（例如，hk.Linear、hk.nets.MLP 等）也是如此。如果您好奇，请查看如果您将 self.inner_simple 分配给 hk.Linear 的实例而不是 SimpleModule 的实例会发生什么。

现在让我们尝试多层嵌套

[ ]:

class TwiceNestedModule(hk.Module):
  """A module class with a nested module containing a nested module."""

  def __init__(self, output_channels, name=None):
    super().__init__(name)
    assert isinstance(output_channels, int)
    self._output_channels = output_channels
    self.inner_nested = NestedModule(self._output_channels)

  def __call__(self, x):
    w_shape = (x.shape[-1], self._output_channels)
    w = hk.get_parameter("w", w_shape, x.dtype, init=w_init)
    return jnp.dot(x, w) + self.inner_nested(x)

[ ]:

def f(x):
  """Create the module instances and inspect their names."""
  # Instantiate a NestedModule instance. This will be named `a_nested_module`.
  # The SimpleModule instance created inside it will be named
  # a_nested_module/a_simple_module`.
  outer = TwiceNestedModule(output_channels=2)
  outer_out = outer(x)
  print(f'The name assigned to the most outer class is: "{outer.module_name}".')
  print('The name assigned to the module inside "double_nested" is: "'
        f'{outer.inner_nested.module_name}".')
  print('The name assigned to the module inside it is "'
        f'{outer.inner_nested.inner_simple.module_name}".')
  return outer_out

transform_and_print_shapes(f)

The name assigned to the most outer class is: "twice_nested_module".
The name assigned to the module inside "double_nested" is: "twice_nested_module/~/nested_module".
The name assigned to the module inside it is "twice_nested_module/~/nested_module/~/simple_module".

The name and shape of the parameters are:
{'twice_nested_module': {'w': (3, 2)}, 'twice_nested_module/~/nested_module': {'w': (3, 2)}, 'twice_nested_module/~/nested_module/~/simple_module': {'w': (3, 2)}}

太棒了，这也按预期工作：模块名称和调用的完整层次结构反映在内部模块名称中。

Haiku 文档

Haiku 中的参数共享

目录

Haiku 文档

Haiku 中的参数共享

目录

Haiku 中的参数共享#

简介#

扁平模块（无嵌套）#

嵌套模块#

多重转换：合并参数但不共享它们#

在转换后的函数之间共享参数#

案例 1：所有模块具有相同的名称和相同的形状#

案例 2：公共模块具有相同的名称和相同的形状#

案例 3：公共模块具有相同的名称，但不同的形状#