Siamese LSTM模型GPU运行报错：RNN权重不连续及形状不匹配

我正在开发用于分类任务的Siamese LSTM模型，CPU上运行正常，但切换到GPU时出现两个问题：

• UserWarning：RNN模块权重不属于单一连续内存块，建议调用flatten_parameters()
• RuntimeError：形状'[512, 1]'对大小为34304的输入无效

已尝试调用flatten_parameters()和让双分支共享同一LSTM实例，问题仍未解决，相关代码及报错堆栈如下：

# Check if CUDA is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Define the Siamese LSTM model
class SiameseLSTM(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, num_classes):
        super(SiameseLSTM, self).__init__()
        self.encoder = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, num_classes)  # Output layer with num_classes units

    def forward(self, x1, x2):
        out1, _ = self.encoder(x1)
        out1 = out1[:, -1, :]  
        out2, _ = self.encoder(x2) 
        out2 = out2[:, -1, :]  
        out1 = F.softmax(self.fc(out1), dim=1)  
        out2 = F.softmax(self.fc(out2), dim=1)  
        return out1, out2

class CSVDataset(Dataset):
    def __init__(self, folder_path, transform=None):
        self.folder_path = folder_path
        self.transform = transform
        self.file_paths = [os.path.join(folder_path, file) for file in os.listdir(folder_path) if file.endswith('.csv')]
        self.labels = [self.extract_label(file) for file in self.file_paths]  # Extract labels from file names

    def __len__(self):
        return len(self.file_paths)

    def extract_label(self, file_name):
        if 'wheat' in file_name:
            return 1
        elif 'mustard' in file_name:
            return 2
        elif 'sugarcane' in file_name:
            return 3
        else:
            return 0  # If none of the keywords are present, assign label 0

    def __getitem__(self, idx):
        data = pd.read_csv(self.file_paths[idx])
        if 'feature_index' in data.columns:
            data.drop(columns=['feature_index'], inplace=True)
        
        if self.transform:
            data = self.transform(data)

        label = self.labels[idx]
        return data, label

    def collate_fn(self, batch):
        padded_batch = [seq.clone().detach() for seq, _ in batch]
        padded_batch = pad_sequence(padded_batch, batch_first=True, padding_value=0.0)
        labels = [label for _, label in batch]
        return padded_batch, labels

# Custom transform function to convert data to PyTorch tensors
def transform_fn(data):
    if 'date' in data.columns:
        data.drop(columns=['date'], inplace=True)
    data_tensor = torch.tensor(data.values, dtype=torch.float32)
    return data_tensor

folder_path = r'E:\project_data\final data'
dataset = CSVDataset(folder_path, transform=transform_fn)

# Split the dataset into training and testing sets
train_indices, test_indices = train_test_split(list(range(len(dataset))), test_size=0.2, random_state=42)

# Create Subset objects for train and test datasets
train_dataset = Subset(dataset, train_indices)
test_dataset = Subset(dataset, test_indices)

# Define data loaders
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, collate_fn=dataset.collate_fn)
test_loader = DataLoader(test_dataset, batch_size=64, collate_fn=dataset.collate_fn)

# Hyperparameters
input_size = len(train_dataset[0][0])  
hidden_size = 128
num_layers = 20
num_classes = len(set(dataset.labels))  

# Model, loss, optimizer
model = SiameseLSTM(input_size, hidden_size, num_layers, num_classes).to(device)
model.encoder.flatten_parameters()  
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Training loop
num_epochs = 10
for epoch in range(num_epochs):
    model.train()  # Set the model to training mode
    running_loss = 0.0
    for data, labels in train_loader:
        data, labels = data.to(device), torch.tensor(labels).to(device)  # Move data to the GPU
        optimizer.zero_grad()
        outputs1, outputs2 = model(data[:, 0, :, None], data[:, 1, :, None])
        
        # Convert output probabilities to class indices
        target1 = torch.argmax(outputs1, dim=1)
        target2 = torch.argmax(outputs2, dim=1)
        
        # Compute cross-entropy loss
        loss = criterion(outputs1, target1) + criterion(outputs2, target2)
        
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch {epoch+1}/{num_epochs}, Loss: {running_loss/len(train_loader)}')

报错信息：

E:\Anaconda\envs\torch\lib\site-packages\torch\nn\modules\rnn.py:878: UserWarning: RNN module weights are not part of single contiguous chunk of memory. This means they need to be compacted at every call, possibly greatly increasing memory usage. To compact weights again call flatten_parameters(). (Triggered internally at C:\cb\pytorch_1000000000000\work\aten\src\ATen\native\cudnn\RNN.cpp:982.)
  result = _VF.lstm(input, hx, self._flat_weights, self.bias, self.num_layers,
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[24], line 106
    104 data, labels = data.to(device), torch.tensor(labels).to(device)  # Move data to the GPU
    105 optimizer.zero_grad()
--> 106 outputs1, outputs2 = model(data[:, 0, :, None], data[:, 1, :, None])
    108 # Convert output probabilities to class indices
    109 target1 = torch.argmax(outputs1, dim=1)

...（中间堆栈省略）

RuntimeError: shape '[512, 1]' is invalid for input of size 34304

问题根源与修复步骤 #

1. RuntimeError：输入维度不匹配 #

你的LSTM输入维度错误，导致内部权重矩阵运算时形状不兼容：

• 训练时传入的data[:, 0, :, None]会让输入形状变成[batch_size, 1, feature_size, 1]，但LSTM（batch_first=True时）要求的格式是[batch_size, seq_len, input_size]
• 多余的维度让LSTM误把特征数当成序列长度，1当成输入特征数，和初始化时的input_size不匹配，引发形状错误

修复代码：
修改训练循环中的模型调用，去掉多余维度：

# 原代码
outputs1, outputs2 = model(data[:, 0, :, None], data[:, 1, :, None])
# 修改后
outputs1, outputs2 = model(data[:, 0:1, :], data[:, 1:2, :])

2. UserWarning：权重内存不连续 #

只在模型初始化后调用一次flatten_parameters()不够，模型移动到GPU后权重布局会变化，需要：

• 模型移到GPU之后调用flatten_parameters()
• 如果用多线程DataLoader，建议在forward方法开头调用，确保每次前向传播前参数连续

修复代码：

# 模型移动到GPU后调用
model = SiameseLSTM(input_size, hidden_size, num_layers, num_classes).to(device)
model.encoder.flatten_parameters()

# 或者在forward方法中添加
def forward(self, x1, x2):
    self.encoder.flatten_parameters()  # 每次前向传播前确保参数连续
    out1, _ = self.encoder(x1)
    # ... 其余代码不变

3. 损失计算逻辑错误 #

你的损失计算完全不符合CrossEntropyLoss的要求：

• CrossEntropyLoss要求输入是未经过softmax的logits，你在模型中提前做了softmax，会导致损失计算异常
• 用torch.argmax(outputs1, dim=1)作为target是错误的，target应该是原始的类别标签，不是模型的预测结果

修复代码：
首先修改模型forward方法，去掉softmax：

def forward(self, x1, x2):
    out1, _ = self.encoder(x1)
    out1 = out1[:, -1, :]  
    out2, _ = self.encoder(x2) 
    out2 = out2[:, -1, :]  
    out1 = self.fc(out1)  # 去掉F.softmax
    out2 = self.fc(out2)  # 去掉F.softmax
    return out1, out2

然后修改损失计算：

# 原代码
target1 = torch.argmax(outputs1, dim=1)
target2 = torch.argmax(outputs2, dim=1)
loss = criterion(outputs1, target1) + criterion(outputs2, target2)

# 修改后（Siamese任务用两个样本输出和同一标签计算损失）
loss = criterion(outputs1, labels) + criterion(outputs2, labels)

4. 额外优化建议 #

• num_layers=20层数过多，极易出现梯度消失，建议调整为2-4层
• DataLoader的collate_fn中，padded_batch = [seq.clone().detach() for seq, _ in batch]不需要clone，直接用原始序列即可

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