在Google Colab Pro上训练PyTorch车道检测模型崩溃求助

一、报错原因并非数据集过大 #

你遇到的PIL.UnidentifiedImageError和数据集规模无关，核心是单张图像加载失败，可能的诱因包括：

• 图像文件实际损坏（肉眼无法识别的文件头损坏、部分字节丢失）
• Google Drive挂载后的文件读取延迟/异常
• 多进程数据加载时的资源冲突

二、解决PIL.UnidentifiedImageError的具体措施 #

• 验证图像真实状态：用代码直接校验目标图像的完整性，而非仅肉眼查看：

  from PIL import Image
  
  img_path = '/content/drive/MyDrive/path_to_dataset/data/train/image-188463.jpg'
  try:
      with Image.open(img_path) as img:
          img.verify()  # 深层验证图像编码
          print("图像格式正常")
          # 尝试重新加载并显示像素
          img.load()
          print(f"图像尺寸: {img.size}")
  except Exception as e:
      print(f"图像存在损坏: {str(e)}")
  

  若确认损坏，直接删除或替换该图像；若代码验证正常，则问题出在数据加载流程。

• 调整多进程加载参数：num_workers=4在读取Drive文件时容易引发多进程竞争，先将num_workers设为0（单进程加载）测试，若恢复正常，再逐步调整为1或2，平衡速度与稳定性。

• 在Dataset中添加异常捕获：修改自定义Dataset的__getitem__方法，跳过损坏图像，避免训练中断：

  def __getitem__(self, idx):
      img_path = self.image_paths[idx]
      try:
          image = Image.open(img_path).convert('RGB')
          # 后续图像预处理、标签读取逻辑
          return image, label
      except (PIL.UnidentifiedImageError, IOError) as e:
          print(f"跳过损坏图像: {img_path}")
          # 随机返回另一张正常图像，保证训练流程不中断
          return self.__getitem__((idx + 1) % len(self.image_paths))
  

三、解决Google Drive到Colab会话的文件传输问题 #

• 用压缩包+gdown批量迁移：先将数据集打包为ZIP文件上传至Drive，设置为可共享，再用gdown下载到Colab本地：

  # 安装gdown工具
  !pip install -q gdown
  # 替换为你的Drive压缩包文件ID（从共享链接提取）
  !gdown --id YOUR_DATASET_ZIP_ID
  # 解压到Colab本地存储（/content/目录读写速度远快于Drive）
  !unzip -q dataset.zip -d /content/local_dataset/
  

  文件ID提取方式：Drive共享链接格式为https://drive.google.com/file/d/XXXXXX/view，其中XXXXXX即为文件ID。

• 挂载Drive后批量复制到本地：若必须保留Drive挂载，先重新挂载并延长超时时间，再用rsync断点续传复制数据到Colab本地：

  from google.colab import drive
  # 延长挂载超时时间，避免大文件挂载失败
  drive.mount('/content/drive', timeout_ms=120000)
  

  # 用rsync断点续传复制训练集到本地
  !rsync -avP /content/drive/MyDrive/path_to_dataset/data/train /content/local_dataset/
  # 复制验证集
  !rsync -avP /content/drive/MyDrive/path_to_dataset/data/val /content/local_dataset/
  

  后续训练直接读取/content/local_dataset/下的数据，既稳定又能提升加载速度。

四、训练稳定性优化 #

• 调整batch_size：batch_size=256可能占用过多GPU显存，导致训练后期内存溢出崩溃，先降至128测试稳定性，再根据显存使用情况逐步调整。

• 梯度累积模拟大batch效果：若想保留大batch的训练效果，用梯度累积替代直接设置大batch：

  accumulation_steps = 2  # 累积2次小batch，等效于batch_size=256
  optimizer.zero_grad()
  
  for step, (images, labels) in enumerate(train_loader):
      outputs = model(images)
      loss = criterion(outputs, labels)
      # 损失除以累积步数，保证梯度尺度正确
      loss = loss / accumulation_steps
      loss.backward()
  
      # 每累积steps步更新一次参数
      if (step + 1) % accumulation_steps == 0:
          optimizer.step()
          optimizer.zero_grad()
  

• 定期保存训练 checkpoint：在训练循环中添加自动保存逻辑，避免崩溃前丢失进度：

  # 每epoch结束保存一次模型
  if epoch % 1 == 0:
      save_path = f'/content/checkpoint_epoch_{epoch}.pth'
      torch.save({
          'epoch': epoch,
          'model_state_dict': model.state_dict(),
          'optimizer_state_dict': optimizer.state_dict(),
          'loss': current_loss,
      }, save_path)
  

内容的提问来源于stack exchange，提问作者Mustafa Uğur Baskın