TensorFlow与PyTorch层规范化/全连接层形状查看及行为对齐问题

问题背景 #

在TensorFlow中创建LayerNormalization和Dense层后，直接打印层对象只会输出实例地址，无法像PyTorch那样直接查看参数形状；同时需要确保两个框架的层行为逻辑一致。

PyTorch示例代码（可正常输出参数形状） #

import torch
import torch.nn as nn

data = [3, 3, 3, 3, 4, 3, 3, 3, 3, 4, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 16, 16, 16, 16, 16, 16, 16, 17]
aux_idxs_torch = torch.tensor(data)
layer_norms_torch = nn.ModuleList([])
linear_layers_torch = nn.ModuleList([])
for i in aux_idxs_torch:
    layer_norms_torch.append(nn.LayerNorm([2*i, 100]))
    linear_layers_torch.append(nn.Linear(2*i, 128))
print(layer_norms_torch[0].weight.shape)
print(linear_layers_torch[0].weight.shape)

输出：

torch.Size([6, 100])
torch.Size([128, 6])

原始TensorFlow代码（仅输出实例地址） #

import tensorflow as tf

data = [3, 3, 3, 3, 4, 3, 3, 3, 3, 4, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 16, 16, 16, 16, 16, 16, 16, 17]
aux_idxs_tf = tf.constant(data)
layer_norms_tf = []
linear_layers_tf = []
for i in aux_idxs_tf:
    layer_norms_tf.append(tf.keras.layers.LayerNormalization(input_shape=(2*i, 100)))
    linear_layers_tf.append(tf.keras.layers.Dense(128, input_shape=(2*i,)))
print(layer_norms_tf[0])
print(linear_layers_tf[0])

输出：

<keras.layers.normalization.layer_normalization.LayerNormalization object at 0x7f91f862abf0>
<keras.layers.core.dense.Dense object at 0x7f91f87eb7f0>

解决方案 #

一、修改TensorFlow代码查看参数形状 #

TensorFlow的Keras层采用延迟构建机制，只有输入数据通过层或显式调用build()后，才会初始化权重参数。有两种方式实现：

方式1：显式调用build()初始化权重

import tensorflow as tf

data = [3, 3, 3, 3, 4, 3, 3, 3, 3, 4, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 16, 16, 16, 16, 16, 16, 16, 17]
aux_idxs_tf = tf.constant(data)
layer_norms_tf = []
linear_layers_tf = []

for i in aux_idxs_tf:
    # 将TensorFlow张量转为Python整数，用于shape参数
    input_dim = int(2 * i)
    
    # 构建LayerNormalization层并初始化参数
    ln_layer = tf.keras.layers.LayerNormalization(input_shape=(input_dim, 100))
    ln_layer.build(input_shape=(None, input_dim, 100))
    layer_norms_tf.append(ln_layer)
    
    # 构建Dense层并初始化参数
    dense_layer = tf.keras.layers.Dense(128, input_shape=(input_dim,))
    dense_layer.build(input_shape=(None, input_dim))
    linear_layers_tf.append(dense_layer)

# 查看参数形状
print(layer_norms_tf[0].gamma.shape)  # LayerNormalization的权重对应gamma
print(linear_layers_tf[0].kernel.shape)  # Dense的权重对应kernel

输出：

(6, 100)
(6, 128)

注意：TensorFlow的Dense层权重形状为(输入维度, 输出维度)，PyTorch的nn.Linear为(输出维度, 输入维度)，但计算逻辑一致（PyTorch是y = x @ weight.T + bias，TensorFlow是y = x @ kernel + bias）。

方式2：传入样例数据触发层构建

import tensorflow as tf

data = [3, 3, 3, 3, 4, 3, 3, 3, 3, 4, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 16, 16, 16, 16, 16, 16, 16, 17]
aux_idxs_tf = tf.constant(data)
layer_norms_tf = []
linear_layers_tf = []

for i in aux_idxs_tf:
    input_dim = int(2 * i)
    
    ln_layer = tf.keras.layers.LayerNormalization(input_shape=(input_dim, 100))
    # 传入样例输入触发层初始化
    ln_layer(tf.random.normal((1, input_dim, 100)))
    layer_norms_tf.append(ln_layer)
    
    dense_layer = tf.keras.layers.Dense(128, input_shape=(input_dim,))
    dense_layer(tf.random.normal((1, input_dim)))
    linear_layers_tf.append(dense_layer)

# 查看参数形状
print(layer_norms_tf[0].gamma.shape)
print(linear_layers_tf[0].kernel.shape)

二、让TensorFlow层列表展示更清晰 #

通过调用层的get_config()方法，可以输出层的关键配置参数，类似PyTorch的层规格展示：

# 打印LayerNormalization层列表详情
for idx, ln_layer in enumerate(layer_norms_tf):
    print(f"LayerNormalization {idx}: {ln_layer.get_config()}")

# 打印Dense层列表详情
for idx, dense_layer in enumerate(linear_layers_tf):
    print(f"Dense {idx}: {dense_layer.get_config()}")

三、验证TensorFlow与PyTorch层行为一致 #

需要对齐初始化方式、归一化维度等关键参数，再通过数值验证确保输出一致：

1. 对齐LayerNormalization行为

PyTorch的nn.LayerNorm默认对传入的所有维度归一化，TensorFlow默认仅对最后一维归一化，需手动设置axis参数：

# 修改LayerNormalization创建代码，对齐归一化维度
ln_layer = tf.keras.layers.LayerNormalization(axis=[-2, -1], input_shape=(input_dim, 100))

2. 对齐Dense层初始化方式

PyTorch的nn.Linear默认用He均匀初始化，TensorFlow默认用Glorot均匀初始化，需手动指定初始化器：

# 修改Dense层创建代码，对齐初始化方式
dense_layer = tf.keras.layers.Dense(128, input_shape=(input_dim,),
                                    kernel_initializer=tf.keras.initializers.HeUniform())

3. 数值验证输出一致性

import torch

# 生成相同输入数据
torch_ln_input = torch.randn(1, 6, 100)
tf_ln_input = tf.convert_to_tensor(torch_ln_input.numpy())

# 对比LayerNormalization输出
torch_ln_out = layer_norms_torch[0](torch_ln_input)
tf_ln_out = layer_norms_tf[0](tf_ln_input)
print(torch.allclose(torch_ln_out, torch.tensor(tf_ln_out.numpy()), atol=1e-5))  # 应输出True

# 对比Dense层输出
torch_dense_input = torch.randn(1, 6)
tf_dense_input = tf.convert_to_tensor(torch_dense_input.numpy())
torch_dense_out = linear_layers_torch[0](torch_dense_input)
tf_dense_out = linear_layers_tf[0](tf_dense_input)
print(torch.allclose(torch_dense_out, torch.tensor(tf_dense_out.numpy()), atol=1e-5))  # 应输出True

内容的提问来源于stack exchange，提问作者nic.o