PyTorch中3DCNN+LSTM训练脚本无预警中断问题排查求助

问题描述 #

运行基于PyTorch的3DCNN+LSTM训练脚本，原脚本仅适配3DCNN，现在训练会在第1或第2个batch后突然中断，无任何报错或警告。已排除system exit语句、内存不足（调整过batch size），训练能完成第一个batch，但第二个batch中模型无法返回输出导致脚本终止。模型架构如下：

def __init__(self, num_classes):
    super(ConvLSTM, self).__init__()
    self.conv_layer1 = self._make_conv_layer(3, 64, (1, 2, 2), (1, 2, 2))
    self.conv_layer2 = self._make_conv_layer(64, 128, (2, 2, 2), (2, 2, 2))
    self.conv_layer3 = self._make_conv_layer(128, 256, (2, 2, 2), (2, 2, 2))
    self.conv_layer4 = self._make_conv_layer(256, 256, (2, 2, 2), (2, 2, 2))

    self.lstm = nn.LSTM(input_size=256, hidden_size=512, num_layers=1, batch_first=True)

    self.fc5 = nn.Linear(512, 512)
    self.fc5_act = nn.ELU()
    self.fc6 = nn.Linear(512, num_classes)

def _make_conv_layer(self, in_c, out_c, pool_size, stride):
    conv_layer = nn.Sequential(
        nn.Conv3d(in_c, out_c, kernel_size=3, stride=1, padding=1),
        nn.BatchNorm3d(out_c),
        nn.ELU(),
        nn.MaxPool3d(pool_size, stride=stride, padding=0)
    )
    return conv_layer

def forward(self, x):
    # print(f"input {x.size()}")
    x = self.conv_layer1(x)
    x = self.conv_layer2(x)
    x = self.conv_layer3(x)
    x = self.conv_layer4(x)
    
    x = x.permute(0, 2, 1, 3, 4).contiguous()
    batch_size, seq_len, input_size, height, width = x.size()
    x = x.view(batch_size, seq_len, -1)
    
    lstm_out, _ = self.lstm(x)
    lstm_out = lstm_out[:, -1, :]

    x = self.fc5(lstm_out)
    x = self.fc5_act(x)
    x = self.fc6(x)
    return x

可能原因与排查步骤 #

1. 3DCNN到LSTM的维度不匹配 #

从代码的forward函数看，3DCNN输出经过permute和view后传入LSTM，但LSTM的input_size硬编码为256。如果x.view(batch_size, seq_len, -1)后的最后一维不等于256，会导致LSTM输入维度不匹配。在CUDA环境下，这种错误可能不会抛出明确Python报错，直接导致进程终止。

• 排查方法：在forward函数的关键节点添加维度打印：

  def forward(self, x):
      print(f"Input shape: {x.size()}")
      x = self.conv_layer1(x)
      print(f"After conv1: {x.size()}")
      x = self.conv_layer2(x)
      print(f"After conv2: {x.size()}")
      x = self.conv_layer3(x)
      print(f"After conv3: {x.size()}")
      x = self.conv_layer4(x)
      print(f"After conv4: {x.size()}")
      
      x = x.permute(0, 2, 1, 3, 4).contiguous()
      print(f"After permute: {x.size()}")
      batch_size, seq_len, input_size, height, width = x.size()
      x = x.view(batch_size, seq_len, -1)
      print(f"Before LSTM: {x.size()}")  # 确认最后一维是否为256
      
      lstm_out, _ = self.lstm(x)
      # ... 其余代码
  

  如果维度不匹配，要么调整LSTM的input_size为实际计算出的维度，要么修改3DCNN的卷积/池化参数（比如调整pool_size或stride），让输出维度符合预期。

2. CUDA上下文异常（无报错崩溃的典型场景） #

当CUDA操作出现非法内存访问、张量维度异常等硬件层面错误时，PyTorch可能不会输出Python层面的报错，直接终止进程。

• 排查方法：
  • 切换到CPU训练：将模型和数据都移到CPU上运行，看是否还会崩溃。如果CPU下正常，说明问题出在CUDA相关逻辑。
  • 开启CUDA异常检测：在脚本开头添加torch.autograd.set_detect_anomaly(True)；或者在运行脚本时添加环境变量CUDA_LAUNCH_BLOCKING=1（Linux下执行CUDA_LAUNCH_BLOCKING=1 python your_script.py，Windows下先设置环境变量再运行），强制同步CUDA操作，让错误信息打印出来。

3. 数据加载器的隐式错误 #

第一个batch正常不代表所有batch数据都没问题，第二个batch的数据可能存在损坏、维度不一致或预处理异常。

• 排查方法：
  • 单独遍历DataLoader，打印每个batch的输入张量维度、数据类型：

    for idx, (data, label) in enumerate(dataloader):
        print(f"Batch {idx}: data shape {data.size()}, label shape {label.size()}")
    

    确认所有batch的维度完全一致。
  • 提取第二个batch的数据，手动喂给模型：如果单独喂第二个batch时崩溃，说明该batch数据有问题，检查对应的数据文件是否损坏，预处理逻辑是否正确（比如视频帧数量是否统一、归一化是否出现异常值）。

4. BatchNorm3d的运行时异常 #

3DCNN中的BatchNorm层在训练时需要计算通道均值和方差，如果某个batch的所有样本在某个通道上的数值完全相同，会导致方差为0，引发除以0的错误。这种错误在CUDA优化下可能不会抛出明确报错，直接终止进程。

• 排查方法：
  • 在BatchNorm层后添加打印，输出每个通道的均值和方差，或者临时注释掉BatchNorm层，看训练是否能继续。
  • 关闭cudnn加速：在脚本开头添加torch.backends.cudnn.enabled = False，排查是否是cudnn优化导致的问题。

5. LSTM隐藏状态初始化问题 #

虽然代码中使用LSTM默认的隐藏状态初始化，但在某些场景下（比如多GPU训练、自定义数据并行），可能存在隐藏状态未正确初始化的问题。

• 排查方法：手动初始化LSTM的隐藏状态，确保每次forward都使用全新的状态：

  def forward(self, x):
      # ... 3DCNN部分代码
      x = x.view(batch_size, seq_len, -1)
      
      # 手动初始化隐藏状态
      h0 = torch.zeros(self.lstm.num_layers, batch_size, self.lstm.hidden_size).to(x.device)
      c0 = torch.zeros(self.lstm.num_layers, batch_size, self.lstm.hidden_size).to(x.device)
      lstm_out, _ = self.lstm(x, (h0, c0))
      # ... 其余代码
  

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