FastAI与PyTorch在Colab训练过慢：GPU配置问题排查求助

我正在用FastAI+PyTorch做图像分类任务，在Google Colab上训练时速度特别慢，已经确认用了GPU运行时，代码如下：

import os
import torch
import torchvision as tv
import matplotlib.pyplot as plt
import numpy as np
from torchinfo import summary as torchinfo_summary

#cuda configs
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

#loading dataset
data_dir='/content/drive/MyDrive/tech_related/machine_learning_related/pytorch-etic/data/cat_deer_dog_horse'
#data_dir = os.path.join('..','data','cat_deer_dog_horse')
print(os.listdir(data_dir))
data_classes = os.listdir(os.path.join(data_dir,'train'))
print(data_classes)

#making torch defined dataloaders with dataset
stats = ((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))  #means and stds of each channel in images of cifar10
train_tfms = tv.transforms.Compose([
    tv.transforms.RandomCrop(32,padding=4,padding_mode='reflect'),
    tv.transforms.RandomHorizontalFlip(),
    tv.transforms.ToTensor()
])
valid_tfms = tv.transforms.Compose([
    tv.transforms.ToTensor()
])

train_ds = tv.datasets.ImageFolder(os.path.join(data_dir,'train'),train_tfms)
valid_ds = tv.datasets.ImageFolder(os.path.join(data_dir,'test'),valid_tfms)
batch_size = 64
train_dl = torch.utils.data.DataLoader(train_ds,batch_size,shuffle=True,pin_memory=True)
valid_dl = torch.utils.data.DataLoader(valid_ds,batch_size,shuffle=True,pin_memory=True)

model_ = tv.models.mobilenet_v2(pretrained=False, num_classes=len(data_classes)).to(device)

#here comes the training part, finding optimum learning rate.
from fastai.vision.all import *
data = DataLoaders(train_dl,valid_dl)
learner = Learner(data, model_, loss_func=torch.nn.functional.cross_entropy, opt_func=Adam, metrics=accuracy)

lr_min,lr_steep,lr_slide,lr_valley = learner.lr_find(suggest_funcs=(minimum,steep,slide,valley))

问题排查与优化方案 #

• 数据集存储位置导致IO瓶颈
  你的数据集存在Google Drive中，Colab读取Drive文件的速度远低于本地临时存储，每次数据加载都会拖慢训练。建议把数据集复制到Colab的临时目录：

  !cp -r /content/drive/MyDrive/tech_related/machine_learning_related/pytorch-etic/data/cat_deer_dog_horse /content/
  data_dir = '/content/cat_deer_dog_horse'
  

• DataLoader未启用多线程加载
  代码中DataLoader的num_workers参数默认是0，仅用主线程加载数据，会导致GPU等待数据闲置。Colab中可设置num_workers=2或4，配合pin_memory=True实现数据加载与GPU计算并行：

  train_dl = torch.utils.data.DataLoader(train_ds,batch_size,shuffle=True,pin_memory=True, num_workers=2)
  valid_dl = torch.utils.data.DataLoader(valid_ds,batch_size,shuffle=False,pin_memory=True, num_workers=2)
  

• 验证集无需打乱
  验证集设置shuffle=True无意义，还会额外消耗资源，直接改为shuffle=False即可。

• 用FastAI内置工具简化数据处理
  既然使用FastAI，推荐直接用ImageDataLoaders.from_folder创建数据加载器，它会自动处理多线程、数据增强优化等细节，比手动用PyTorch DataLoader更高效：

  from fastai.vision.all import *
  data = ImageDataLoaders.from_folder(data_dir, train='train', valid='test', 
                                      item_tfms=Resize(32),
                                      batch_tfms=[RandomCrop(32, padding=4), RandomFlip()])
  learner = Learner(data, mobilenet_v2(pretrained=False, num_classes=4), loss_func=CrossEntropyLossFlat(), metrics=accuracy)
  

• 确认GPU实际运行状态
  即使选择了GPU运行时，也可以手动验证：执行!nvidia-smi查看GPU占用情况，或者在代码中添加print(torch.cuda.is_available(), torch.cuda.get_device_name(0))，确保模型和数据确实在GPU上运行。

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