PyTorch与TensorFlow 1.x GRU层输出不一致问题及解决问询

问题背景 #

我正尝试将TensorFlow 1.x中训练好的多单元GRU模型迁移到PyTorch，以便将编码器与其他PyTorch高级模块结合。已提取各层权重到NumPy数组，在PyTorch中重建模型并手动设置权重后，测试输出的张量维度正确，但数值差异极大。

原始TensorFlow 1.x模型代码 #

encoder_cell = [tf.nn.rnn_cell.GRUCell(size) for size in self.cell_size]
encoder_cell = tf.contrib.rnn.MultiRNNCell(encoder_cell)
encoder_outputs, encoder_state = tf.nn.dynamic_rnn(encoder_cell,
                                                   encoder_emb_inp,
                                                   sequence_length=self.input_len,
                                                   dtype=tf.float32,
                                                   time_major=False)
emb = tf.layers.dense(tf.concat(encoder_state, axis=1),
                              self.embedding_size
                             )
emb = self.emb_activation(emb)

return emb

重建的PyTorch模型代码 #

class GRU_Encoder(nn.Module):
    def __init__(self, input_size=40, embedding_size=32, hidden_size=512):
        super(GRU_Encoder, self).__init__()
        
        self.weights_dir = '/media/drives/drive1/robin/cddd/default_model/weights'
        self.embedding_sizes = [32, 512, 1024, 2048]
        
        # Initial embedding layer
        char_weights = np.load(f'{self.weights_dir}/char_embedding_0.npy')
        self.char_projection = nn.Linear(input_size, embedding_size, bias=False)
        self.char_projection.weight.data = torch.FloatTensor(char_weights.T)
        
        # Create 3 GRU layers
        self.gru_layers = nn.ModuleList([
            nn.GRU(
                input_size=self.embedding_sizes[i],
                hidden_size=self.embedding_sizes[i+1],
                batch_first=True
            ) for i in range(3)
        ])
        
        # Final dense layer
        dense_kernel = np.load(f'{self.weights_dir}/Encoder_dense_kernel_0.npy')
        self.dense = nn.Linear(sum(self.embedding_sizes[1:]), dense_kernel.shape[1])
        self.tanh = nn.Tanh()
        
        self._load_gru_weights()
        self._load_dense_weights()
    
    def _load_gru_weights(self):
        for i in range(3):
            # Load weights
            gates_kernel = np.load(f'{self.weights_dir}/Encoder_rnn_multi_rnn_cell_cell_{i}_gru_cell_gates_kernel_0.npy')
            gates_bias = np.load(f'{self.weights_dir}/Encoder_rnn_multi_rnn_cell_cell_{i}_gru_cell_gates_bias_0.npy')
            candidate_kernel = np.load(f'{self.weights_dir}/Encoder_rnn_multi_rnn_cell_cell_{i}_gru_cell_candidate_kernel_0.npy')
            candidate_bias = np.load(f'{self.weights_dir}/Encoder_rnn_multi_rnn_cell_cell_{i}_gru_cell_candidate_bias_0.npy')
            
            input_size = self.embedding_sizes[i]
            hidden_size = self.embedding_sizes[i+1]
            
            # Properly reshape weights for PyTorch GRU format
            # PyTorch expects (3 * hidden_size, input_size) for ih weights
            # and (3 * hidden_size, hidden_size) for hh weights
            
            # Split input and hidden weights
            gates_kernel_i = gates_kernel[:input_size, :]  # Input weights
            print(gates_kernel_i.shape)
            gates_kernel_h = gates_kernel[input_size:, :]  # Hidden weights
            print(gates_kernel_h.shape)
            candidate_kernel_i = candidate_kernel[:input_size, :]
            candidate_kernel_h = candidate_kernel[input_size:, :]
            
            # Combine weights in PyTorch's expected format
            w_ih = np.concatenate([
                gates_kernel_i[:, :hidden_size],     # reset gate
                gates_kernel_i[:, hidden_size:],     # update gate
                candidate_kernel_i
            ], axis=1)
            
            w_hh = np.concatenate([
                gates_kernel_h[:, :hidden_size],     # reset gate
                gates_kernel_h[:, hidden_size:],     # update gate
                candidate_kernel_h
            ], axis=1)
            
            # Combine biases
            b_ih = np.concatenate([
                gates_bias[:hidden_size],            # reset gate
                gates_bias[hidden_size:],            # update gate
                candidate_bias
            ])
            
            print('You are here')
            print(self.gru_layers[i].bias_ih_l0.shape)
            print(self.gru_layers[i].bias_hh_l0.shape)
            # Set weights and biases
            self.gru_layers[i].weight_ih_l0.data = torch.FloatTensor(w_ih.T)
            self.gru_layers[i].weight_hh_l0.data = torch.FloatTensor(w_hh.T)
            self.gru_layers[i].bias_ih_l0.data = torch.FloatTensor(b_ih)
            self.gru_layers[i].bias_hh_l0.data = torch.zeros_like(self.gru_layers[i].bias_hh_l0)
    
    def _load_dense_weights(self):
        dense_kernel = np.load(f'{self.weights_dir}/Encoder_dense_kernel_0.npy')
        dense_bias = np.load(f'{self.weights_dir}/Encoder_dense_bias_0.npy')
        self.dense.weight.data = torch.FloatTensor(dense_kernel.T)
        self.dense.bias.data = torch.FloatTensor(dense_bias)
    
    def forward(self, x):
        # Initial projection
        x = self.char_projection(x)
        # Process through GRU layers
        hidden_states = []
        current_input = x
        
        for gru in self.gru_layers:
            output, hidden = gru(current_input)
            print(f' Output: {output.shape}')
            print(f' Hidden: {hidden.shape}')
            output_array = np.array(output)
            print(np.array(hidden[-1]))
            hidden_states.append(hidden[-1])
            current_input = output
        
        # Combine hidden states and apply final transformation
        combined_hidden = torch.cat(hidden_states, dim=1)
        x = self.dense(combined_hidden)
        x = self.tanh(x)
        return x

问题细节 #

经过self.char_projection后的x与原始TensorFlow模型的encoder_emb_inp完全一致，但后续输出差异明显。单GRU层测试发现两者输出相似但数值差异较大，最终定位到候选隐藏状态$\hat{h}$计算步骤中的矩阵乘法顺序不同：TensorFlow 1.x采用原始实现，而PyTorch不支持切换该逻辑。

差异根源 #

TensorFlow 1.x的GRUCell与PyTorch的GRU在候选隐藏状态$\hat{h}$的计算逻辑上存在矩阵乘法顺序的核心差异：

• TensorFlow 1.x中，$\hat{h} = \tanh(W_{xh} \cdot x_t + b_x + W_{hh} \cdot (r_t \odot h_{t-1}) + b_h)$，权重矩阵为[input_size, hidden_size]格式，计算时是输入向量左乘权重矩阵。
• PyTorch的GRU中，$\hat{h} = \tanh(W_{ih} x_t + b_{ih} + W_{hh} (r_t \odot h_{t-1}) + b_{hh})$，其中W_ih为[3*hidden_size, input_size]格式，计算时是权重矩阵左乘输入向量（本质是输入向量转置后右乘权重矩阵转置）。

这种矩阵乘法顺序的差异，若权重转置逻辑处理不当，会直接引发数值偏差。

解决方法 #

要让PyTorch输出与TensorFlow一致，需手动实现GRU的前向逻辑，完全对齐TensorFlow的计算步骤，而非使用PyTorch内置的nn.GRU。

1. 自定义GRU层类 #

实现单个GRU单元的前向计算，严格遵循TensorFlow 1.x的GRU公式：

class CustomGRULayer(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        
        # 定义TensorFlow风格的权重（input_size -> hidden_size）
        self.W_gates = nn.Parameter(torch.empty(input_size + hidden_size, 2 * hidden_size))
        self.b_gates = nn.Parameter(torch.empty(2 * hidden_size))
        self.W_candidate = nn.Parameter(torch.empty(input_size + hidden_size, hidden_size))
        self.b_candidate = nn.Parameter(torch.empty(hidden_size))
        
    def forward(self, x, h_prev):
        # x shape: [batch_size, seq_len, input_size]
        # h_prev shape: [batch_size, hidden_size]
        
        seq_len = x.shape[1]
        outputs = []
        
        for t in range(seq_len):
            x_t = x[:, t, :]  # [batch_size, input_size]
            # 拼接输入与上一时刻隐藏状态
            concat = torch.cat([x_t, h_prev], dim=1)  # [batch_size, input_size + hidden_size]
            
            # 计算重置门r_t和更新门z_t
            gates = torch.matmul(concat, self.W_gates) + self.b_gates
            r_t = torch.sigmoid(gates[:, :self.hidden_size])
            z_t = torch.sigmoid(gates[:, self.hidden_size:])
            
            # 计算候选隐藏状态hat_h
            concat_candidate = torch.cat([x_t, r_t * h_prev], dim=1)
            hat_h = torch.tanh(torch.matmul(concat_candidate, self.W_candidate) + self.b_candidate)
            
            # 更新当前隐藏状态h_t
            h_t = (1 - z_t) * h_prev + z_t * hat_h
            outputs.append(h_t.unsqueeze(1))
            h_prev = h_t
        
        outputs = torch.cat(outputs, dim=1)  # [batch_size, seq_len, hidden_size]
        return outputs, h_t.unsqueeze(0)  # 输出格式对齐PyTorch的nn.GRU（hidden是[num_layers, batch_size, hidden_size]）

2. 修改编码器模型的权重加载逻辑 #

在GRU_Encoder中替换nn.GRU为自定义的CustomGRULayer，直接加载TensorFlow导出的权重，无需转置：

class GRU_Encoder(nn.Module):
    def __init__(self, input_size=40, embedding_size=32, hidden_size=512):
        super(GRU_Encoder, self).__init__()
        
        self.weights_dir = '/media/drives/drive1/robin/cddd/default_model/weights'
        self.embedding_sizes = [32, 512, 1024, 2048]
        
        # Initial embedding layer
        char_weights = np.load(f'{self.weights_dir}/char_embedding_0.npy')
        self.char_projection = nn.Linear(input_size, embedding_size, bias=False)
        self.char_projection.weight.data = torch.FloatTensor(char_weights.T)
        
        # 创建3个自定义GRU层
        self.gru_layers = nn.ModuleList([
            CustomGRULayer(
                input_size=self.embedding_sizes[i],
                hidden_size=self.embedding_sizes[i+1]
            ) for i in range(3)
        ])
        
        # Final dense layer
        dense_kernel = np.load(f'{self.weights_dir}/Encoder_dense_kernel_0.npy')
        self.dense = nn.Linear(sum(self.embedding_sizes[1:]), dense_kernel.shape[1])
        self.tanh = nn.Tanh()
        
        self._load_gru_weights()
        self._load_dense_weights()
    
    def _load_gru_weights(self):
        for i in range(3):
            # 加载TensorFlow导出的权重
            gates_kernel = np.load(f'{self.weights_dir}/Encoder_rnn_multi_rnn_cell_cell_{i}_gru_cell_gates_kernel_0.npy')
            gates_bias = np.load(f'{self.weights_dir}/Encoder_rnn_multi_rnn_cell_cell_{i}_gru_cell_gates_bias_0.npy')
            candidate_kernel = np.load(f'{self.weights_dir}/Encoder_rnn_multi_rnn_cell_cell_{i}_gru_cell_candidate_kernel_0.npy')
            candidate_bias = np.load(f'{self.weights_dir}/Encoder_rnn_multi_rnn_cell_cell_{i}_gru_cell_candidate_bias_0.npy')
            
            # 直接赋值，无需转置，完全对齐TensorFlow的权重格式
            self.gru_layers[i].W_gates.data = torch.FloatTensor(gates_kernel)
            self.gru_layers[i].b_gates.data = torch.FloatTensor(gates_bias)
            self.gru_layers[i].W_candidate.data = torch.FloatTensor(candidate_kernel)
            self.gru_layers[i].b_candidate.data = torch.FloatTensor(candidate_bias)
    
    def _load_dense_weights(self):
        dense_kernel = np.load(f'{self.weights_dir}/Encoder_dense_kernel_0.npy')
        dense_bias = np.load(f'{self.weights_dir}/Encoder_dense_bias_0.npy')
        self.dense.weight.data = torch.FloatTensor(dense_kernel.T)
        self.dense.bias.data = torch.FloatTensor(dense_bias)
    
    def forward(self, x):
        # Initial projection
        x = self.char_projection(x)
        # Process through GRU layers
        hidden_states = []
        current_input = x
        h_prev = None
        
        for gru in self.gru_layers:
            if h_prev is None:
                # 初始隐藏状态设为全0，对齐TensorFlow的dynamic_rnn默认行为
                h_prev = torch.zeros(x.shape[0], gru.hidden_size, device=x.device)
            output, h_t = gru(current_input, h_prev)
            hidden_states.append(h_t.squeeze(0))
            current_input = output
            h_prev = h_t.squeeze(0)
        
        # Combine hidden states and apply final transformation
        combined_hidden = torch.cat(hidden_states, dim=1)
        x = self.dense(combined_hidden)
        x = self.tanh(x)
        return x

3. 额外注意事项 #

• 对齐sequence_length处理逻辑：TensorFlow的dynamic_rnn会根据sequence_length截断后续计算并复制最后有效状态，自定义GRU层也需实现这一逻辑（若原始模型使用了该参数）。
• 验证数值精度：使用相同输入张量，对比PyTorch自定义模型与TensorFlow模型的输出，确保误差在浮点精度范围内（如1e-6以下）。

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