hk.LayerNorm 模块介绍

hk.LayerNorm 是 Haiku 库中用于实现 Layer Normalization（层归一化）的模块。Layer Normalization 是一种神经网络归一化的技术，旨在提高神经网络的训练稳定性和泛化性。

主要参数：

• axis（默认为-1）： 沿着哪个轴进行归一化。通常选择最后一个轴，对输入的特征进行归一化。

• create_scale（默认为True）： 是否创建可学习的缩放参数。如果为 True，则会创建一个可学习的缩放参数，用于调整归一化后的值的幅度。

• create_offset（默认为True）： 是否创建可学习的偏置参数。如果为 True，则会创建一个可学习的偏置参数，用于调整归一化后的值的偏移。

• epsilon（默认为1e-5）： 一个小的正数，用于防止除以零的情况。

import haiku as hk
import jax
import jax.numpy as jnp
import pickle

### 自定义LayerNorm模块
class LayerNorm(hk.LayerNorm):
  """LayerNorm module.

  Equivalent to hk.LayerNorm but with different parameter shapes: they are
  always vectors rather than possibly higher-rank tensors. This makes it easier
  to change the layout whilst keep the model weight-compatible.
  """

  def __init__(self,
               axis,
               create_scale: bool,
               create_offset: bool,
               eps: float = 1e-5,
               scale_init=None,
               offset_init=None,
               use_fast_variance: bool = False,
               name=None,
               param_axis=None):
    super().__init__(
        axis=axis,
        create_scale=False,
        create_offset=False,
        eps=eps,
        scale_init=None,
        offset_init=None,
        use_fast_variance=use_fast_variance,
        name=name,
        param_axis=param_axis)
    self._temp_create_scale = create_scale
    self._temp_create_offset = create_offset
     
    #self.scale_init = hk.initializers.Constant(1)
    #self.offset_init = hk.initializers.Constant(0)
    
  def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
    is_bf16 = (x.dtype == jnp.bfloat16)
    if is_bf16:
      x = x.astype(jnp.float32)

    param_axis = self.param_axis[0] if self.param_axis else -1
    param_shape = (x.shape[param_axis],)

    param_broadcast_shape = [1] * x.ndim
    param_broadcast_shape[param_axis] = x.shape[param_axis]
    scale = None
    offset = None
    
    # scale，offset张量的形状必须可扩展到输入数据的形状。
    # 没有显式指定 self.scale_init，self.offset_init参数，
    # 则默认使用 Haiku 库中的默认初始化方法。同 def __init__（）中注释的显式指定
    if self._temp_create_scale:
      scale = hk.get_parameter(
          'scale', param_shape, x.dtype, init=self.scale_init)
      scale = scale.reshape(param_broadcast_shape)

    if self._temp_create_offset:
      offset = hk.get_parameter(
          'offset', param_shape, x.dtype, init=self.offset_init)
      offset = offset.reshape(param_broadcast_shape)

    out = super().__call__(x, scale=scale, offset=offset)

    if is_bf16:
      out = out.astype(jnp.bfloat16)

    return out


with open("Human_HBB_tensor_dict_ensembled.pkl",'rb') as f:
  Human_HBB_tensor_dict = pickle.load(f)

input_data = jnp.array(Human_HBB_tensor_dict['msa_feat'])
print(input_data.shape)

# 转换为Haiku模块
# LayerNorm层，在数据最后一个维度/轴（特征）做归一化，并创建可学习的缩放参数和偏置参数
model = hk.transform(lambda x: LayerNorm(axis=[-1], 
                                         create_scale=True,
                                         create_offset=True,
                                         name='msa_feat_norm')(x))
 
print(model)
                                     
## 获取初始化的参数,参数的形状需要输入数据的形状以及模型的结构
rng = jax.random.PRNGKey(42)
params = model.init(rng, input_data)
print(params) 
print("params scale shape:") 
#print(params['msa_feat_norm']['scale'].shape)
#print("params offset bias:")
#print(params['msa_feat_norm']['offset'].shape)
 
output_data = model.apply(params, rng, input_data)
print("input_data shape:", input_data.shape) 
print("Output Data shape:", output_data.shape)
#print("原始数据:", input_data)
print("经过LayerNorm后:", output_data)

### 使用原始的hk.LayerNorm模块
model2 = hk.transform(lambda x: hk.LayerNorm(axis=[-1], 
                                            create_scale=True,
                                            create_offset=True,
                                            name='msa_feat_norm')(x))
 
print(model2)
                                     
params2 = model2.init(rng, input_data)
print(params2) 
print("params2 scale shape:") 
print(params2['msa_feat_norm']['scale'].shape)
print("params2 offset bias:")
print(params2['msa_feat_norm']['offset'].shape)
 
output_data2 = model2.apply(params2, rng, input_data)
print("input_data shape:", input_data.shape) 
print("Output Data shape:", output_data2.shape)
#print("原始数据:", input_data)
print("经过LayerNorm后:", output_data2)

参考：

https://dm-haiku.readthedocs.io/en/latest/api.html?highlight=layernorm#layernorm