TensorFlow模型迁移至PyTorch训练失效问题排查求助

我在TensorFlow中训练了一个二元分类模型，任务为根据描述象限的文本输入与(x, y, z)坐标，预测坐标的(x, z)部分是否属于该象限。该TensorFlow模型准确率接近100%，但将其复现到PyTorch后，模型性能仅为随机水平。我已尝试导入TensorFlow模型权重至对应层，问题仍未解决，请求排查。

TensorFlow实现代码 #

x_data_placeholder = tf.placeholder(tf.float32, [FLAGS.batch_size, 1], name="x_data")
y_data_placeholder = tf.placeholder(tf.float32, [FLAGS.batch_size, 1], name="y_data")
z_data_placeholder = tf.placeholder(tf.float32, [FLAGS.batch_size, 1], name="z_data")
# text and labels placeholders
text_data = tf.placeholder(tf.int32, [FLAGS.batch_size, maxtextlength])
text_lengths = tf.placeholder(tf.int32, [FLAGS.batch_size])
y_labels_placeholder = tf.placeholder(tf.int64, [FLAGS.batch_size])

# encode text and coordinate
embeddings = tf.Variable(tf.random_uniform([100, embedding_size], -1, -1))
rnn_inputs = tf.nn.embedding_lookup(embeddings, text_data)
rnn_layers = [tf.compat.v1.nn.rnn_cell.LSTMCell(size, initializer=tf.compat.v1.keras.initializers.glorot_normal) for size in [256]]
multi_rnn_cell = tf.compat.v1.nn.rnn_cell.MultiRNNCell(rnn_layers, state_is_tuple=True)

text_outputs, text_fstate = tf.compat.v1.nn.dynamic_rnn(cell=multi_rnn_cell,
                                             inputs=rnn_inputs,
                                             dtype=tf.float32, sequence_length=text_lengths)

# have fully connected layers to map them the input coordinates into the same dimension as the LSTM output layer from above
x_output_layer = tf.compat.v1.layers.dense(x_data_placeholder, units=FLAGS.fc_column_size, activation=tf.nn.relu, name='x_coordinate')
y_output_layer = tf.compat.v1.layers.dense(y_data_placeholder, units=FLAGS.fc_column_size, activation=tf.nn.relu, name='y_coordinate')
z_output_layer = tf.compat.v1.layers.dense(z_data_placeholder, units=FLAGS.fc_column_size, activation=tf.nn.relu, name='z_coordinate')

# add the representations
total_output_layer = x_output_layer + y_output_layer + z_output_layer + lstm_output_layer

# make the predictions with two fully connected layers
fc_1 = tf.compat.v1.layers.dense(total_output_layer, units=FLAGS.hidden_layer_size, activation=tf.nn.relu, name='fc_1')
logits = tf.compat.v1.layers.dense(fc_1, units=FLAGS.output_dims, activation=None, name='logits')

# train the model
loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y_labels_placeholder, logits=logits))
optimizer = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate, epsilon=1e-7)
gradients, variables = zip(*optimizer.compute_gradients(loss))
gradients, _ = tf.clip_by_global_norm(gradients, FLAGS.gradient_clip_threshold)
optimize = optimizer.apply_gradients(zip(gradients, variables))

# then it'll be trained with sess.run ...

PyTorch复现代码 #

class BaselineModel(nn.Module):
    def __init__(self):
        super(BaselineModel, self).__init__()

        self.encode_x = nn.Linear(1, embed_size)
        self.encode_y = nn.Linear(1, embed_size)
        self.encode_z = nn.Linear(1, embed_size)
        self._embeddings = nn.Embedding(vocab_size, self.embedding_table_size)
        nn.init.uniform_(self._embeddings.weight, -1.0, 1.0)
        self.num_layers = 1
        self.rnn = nn.LSTM(self.embedding_table_size, self.hidden_size, batch_first=True)
        self.fc_after_text_lstm = nn.Linear(self.hidden_size, 100)
        self.fc = nn.Linear(100, 256) 
        self.fc_final = nn.Linear(256, 2) 
        self.relu_activation = nn.ReLU()
        self.softmax = nn.Softmax(dim=1)

    def init_hidden(self, batch_size, device='cuda:0'):
        # for LSTM, we need # of layers
        h_0 = torch.zeros(1, batch_size, self.hidden_size).to(device)
        c_0 = torch.zeros(1, batch_size, self.hidden_size).to(device)
        return h_0, c_0

    def forward(self, input_text, x_coordinate=None, y_coordinate=None, z_coordinate=None):
        x_embed = self.relu_activation(self.encode_x(x_coordinate.cuda().to(torch.float32)).cuda())
        y_embed = self.relu_activation(self.encode_y(y_coordinate.cuda().to(torch.float32))).cuda()
        z_embed = self.relu_activation(self.encode_z(z_coordinate.cuda().to(torch.float32))).cuda()

        embeds = self._embeddings(input_text)
        embedding, hidden = self.rnn(embeds, self.hidden)
        text_fc = self.relu_activation(self.fc_after_text_lstm(embedding[:, -1]))
        representations_so_far_added = torch.sum(torch.stack([text_fc, x_embed, y_embed, z_embed]), dim=0)
        pre_final_embedding = self.relu_activation(self.fc(representations_so_far_added))
        
        return self.fc_final(pre_final_embedding )

### training code
optimizer = torch.optim.Adam(model.parameters(), lr=0.001, eps=1e-7)
criterion = nn.CrossEntropyLoss()
for input_text, x_coordinate, y_coordinate, z_coordinate, targets, train_data:

    optimizer.zero_grad()
    pred = model(input_text, x_coordinate=x_coordinate, y_coordinate=y_coordinate, z_coordinate=z_coordinate)
    loss = criterion(pred.float(), targets)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
    optimizer.step()
    scheduler.step()

    # accuracy evaluation code, this is evaluated over the entire epoch
    pred_idx = F.log_softmax(pred, dim=1)
    target_labels = targets.cpu().int()
    pred_labels = torch.argmax(pred_idx, dim=-1).cpu().data.int()
    curr_acc = skm.accuracy_score(target_labels, pred_labels)

最小可复现示例 #

import torch
import torch.nn as nn
import numpy as np
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

lr = 0.0005
n_epochs = 10
input_dim = 4    
hidden_dim = 128
layer_dim = 2
output_dim = 2
batch_size = 50

class FeatureDataSet(torch.utils.data.Dataset):

    def __init__(self, x_train, y_train, x_coordinates):
        self.x_train = torch.tensor(x_train, dtype=torch.long)
        self.y_train = torch.tensor(y_train)
        self.x_coordinates = torch.tensor(x_coordinates, dtype=torch.float32)
    def __len__(self):
        return len(self.y_train)
    def __getitem__(self, idx):
        return self.x_train[idx], self.y_train[idx], self.x_coordinates[idx]


class RNN(nn.Module):

    def __init__(self, input_dim, hidden_dim, layer_dim, output_dim, batch_size):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.layer_dim = layer_dim

        # linear layer to encode the coordinate
        self.encode_x = nn.Linear(1, hidden_dim).cuda()
        self._embeddings = nn.Embedding(40, 100).cuda()

        # hidden_dim is 128
        # layer_dim is 2
        self.lstm = nn.LSTM(100, hidden_dim, layer_dim, batch_first=True).cuda()
        self.fc = nn.Linear(2 * hidden_dim, output_dim).cuda()
        self.batch_size = batch_size
        self.hidden = None

    def init_hidden(self, x):
        h0 = torch.zeros(self.layer_dim, x.size(0), self.hidden_dim)
        c0 = torch.zeros(self.layer_dim, x.size(0), self.hidden_dim)
        return [t.cpu() for t in (h0, c0)]

    def forward(self, x, x_coordinate):
        #initializing the hidden states
        h0, c0 = self.init_hidden(x)
        embeds = self._embeddings(x)
        out, (hn, cn) = self.lstm(embeds.cuda(), (h0.cuda(), c0.cuda()))

        x_embed = F.relu(self.encode_x(x_coordinate.cuda().to(torch.float32)).cuda())
        representations_so_far_added = torch.cat([out[:, -1, :], x_embed], dim=1)

        out = self.fc(representations_so_far_added)
        return out

    

model = RNN(input_dim, hidden_dim, layer_dim, output_dim, batch_size)
criterion = nn.CrossEntropyLoss()
opt = torch.optim.Adam(model.parameters(), lr=0.001)

print('Start model training')
import sklearn.metrics as skm
import torch.nn.functional as F


x_train = []
x_coordinates = []
y_train = []
for i in range(10000):
    # create the data. if x_coordinate > 0 and the sentence says that (represented by [1, 5, 6, 8]), then we should predict positive else negative (if the x_coordinate > 0)
    # same applies if the x_coordinate < 0, just that the sentence is now [1, 5, 6, 9]
    if np.random.randint(0, 2) == 0:
        if np.random.randint(0, 2) == 0:
            # x coordinate > 0
            x_train.append([1, 5, 6, 8])
            x_coordinates.append([round(np.random.uniform(0.01, 1.00, 1)[0], 2)])
            y_train.append(1.0)
        else:
            # x coordinate > 0 negative
            x_train.append([1, 5, 6, 8])
            x_coordinates.append([round(np.random.uniform(-1.00, 0.00, 1)[0], 2)])
            y_train.append(0.0)
    else:
        if np.random.randint(0, 2) == 0:
            # x coordinate < 0
            x_train.append([1, 5, 6, 9])
            x_coordinates.append([round(np.random.uniform(-1.00, 0.00, 1)[0], 2)])
            y_train.append(1.0)
        else:
            # x coordinate < 0 negative
            x_train.append([1, 5, 6, 9])
            x_coordinates.append([round(np.random.uniform(0.01, 1.00, 1)[0], 2)])
            y_train.append(0.0)

# print a sample of data 
print(x_train[:10])
print(y_train[:10])
print(x_coordinates[:10])

# create a dataloader
trainingDataset = FeatureDataSet(x_train=x_train, y_train=y_train, x_coordinates=x_coordinates)
train_loader = torch.utils.data.DataLoader(dataset=trainingDataset, batch_size=batch_size, shuffle=True)


# for each epoch
for epoch in range(1, n_epochs + 1):
    acc_all = []
    # each batch
    for i, (x_batch, y_batch, x_coord_batch) in enumerate(train_loader):
        x_batch = x_batch.to(device)
        y_batch = y_batch.to(device)

        x_coord_batch = x_coord_batch.to(device)

        opt.zero_grad()
        # pass in the text (x_batch) and coordinate (x_coord_batch)
        out = model(x_batch, x_coordinate=x_coord_batch)
        loss = criterion(out.float(), y_batch.type(torch.LongTensor).cuda())
        loss.backward()
        opt.step()

        pred_idx = F.log_softmax(out, dim=1)
        target_labels = y_batch.cpu().int()
        pred_labels = torch.argmax(pred_idx, dim=-1).cpu().data.int()

        curr_acc = skm.accuracy_score(target_labels, pred_labels)
        acc_all.append(curr_acc)

    print(np.mean(acc_all))

问题排查与修复方案 #

1. LSTM隐藏状态初始化错误 #

• 问题：PyTorch代码中，BaselineModel的forward方法直接使用未初始化的self.hidden作为LSTM的初始状态；最小可复现示例中init_hidden返回CPU张量后再转CUDA，且未与模型设备保持一致。
• 修复：
  • 在forward方法中每次调用init_hidden初始化当前batch对应的隐藏状态：

    def forward(self, input_text, x_coordinate=None, y_coordinate=None, z_coordinate=None):
        # 初始化当前batch对应的隐藏状态
        self.hidden = self.init_hidden(input_text.size(0), device=input_text.device)
        # ... 后续代码
    

  • init_hidden直接在指定设备创建张量：

    def init_hidden(self, batch_size, device):
        h_0 = torch.zeros(self.num_layers, batch_size, self.hidden_size).to(device)
        c_0 = torch.zeros(self.num_layers, batch_size, self.hidden_size).to(device)
        return (h_0, c_0)
    

2. 特征融合逻辑错误 #

• 问题：TensorFlow中是将文本特征与坐标特征做元素级相加，但PyTorch代码存在维度不匹配和错误的融合方式：
  • BaselineModel中fc_after_text_lstm将LSTM输出转为100维，而坐标编码层encode_x/y/z的输出维度是embed_size，若embed_size≠100会导致维度不匹配；
  • 最小可复现示例中错误使用torch.cat拼接特征，而非TensorFlow中的元素相加。
• 修复：
  • 确保文本特征与坐标特征维度一致：将fc_after_text_lstm的输出维度改为与坐标编码层相同（如256），或调整坐标编码层维度与文本特征一致；
  • 正确实现元素级相加：

    # 替换错误的stack+sum
    representations_so_far_added = text_fc + x_embed + y_embed + z_embed
    

3. 设备与数据类型不统一 #

• 问题：
  • 模型层在__init__中单独调用.cuda()，导致CPU环境下运行报错，且无法灵活切换设备；
  • 训练时目标张量类型与设备未和模型输出统一，如y_batch.type(torch.LongTensor).cuda()未考虑CPU场景。
• 修复：
  • 模型实例化后统一调用.to(device)，移除层初始化时的.cuda()：

    model = BaselineModel().to(device)
    

  • 统一目标张量的类型与设备：

    y_batch = y_batch.long().to(device)
    loss = criterion(out, y_batch)
    

4. LSTM初始化差异 #

• 问题：TensorFlow中LSTM使用glorot_normal初始化，而PyTorch默认初始化方式不同，可能导致模型收敛困难。
• 修复：手动为LSTM权重设置Glorot正态初始化：

  def init_weights(self):
    for name, param in self.rnn.named_parameters():
        if 'weight_ih' in name:
            nn.init.xavier_normal_(param.data)
        elif 'weight_hh' in name:
            nn.init.xavier_normal_(param.data)
        elif 'bias' in name:
            param.data.fill_(0)
  # 在__init__中调用
  self.init_weights()
  

5. 训练循环逻辑错误 #

• 问题：PyTorch训练代码中的循环遍历语法错误（for input_text, x_coordinate, y_coordinate, z_coordinate, targets, train_data:），且未正确处理数据加载。
• 修复：正确遍历DataLoader：

  for batch in train_loader:
      input_text, x_coordinate, y_coordinate, z_coordinate, targets = batch
      # ... 训练逻辑
  

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