模型训练时CPU和GPU大幅度波动——可能是数据的读入拖后腿

模型训练时CPU和GPU大幅度波动——可能是数据的加载拖后腿

问题

在进行猫狗大战分类任务时，发现模型训练时CPU和GPU大幅度波动，且模型训练速度很慢。

原因

​ 初步分析可能是数据加载（包括数据的transform，我用了Resize，ToTensor，Normalize这三个操作）的的速度太慢，于是通过计算一个epoch数据加载的时间来判断，最后发现数据加载的数据和一个epoch训练的时间差不太多(因为用的模型较小，是ResNet18，如果模型比较大，训练时间比数据加载时间大得多的时候，这种情况CPU和GPU的波动频率和幅度会好很多，情况最好的是，在训练一个完batch的数据前，下一个batch的数据已经准备好了)。测量加载时间代码如下：

import time
from torch.utils.data import DataLoader

data_loader = DataLoader(dataset, batch_size=64)
start_time = time.time()

# 遍历数据加载器中的所有批次
for i, data in enumerate(data_loader):
    pass

end_time = time.time()
# 计算并打印整个数据读取的时间
total_time = end_time - start_time
print(f"Total data loading time: {total_time:.4f} seconds")

然后再计算训练一个epoch的时间，若没有比加载数据的时间大很多的话，大概率就是数据加载拖后腿了。

解决方法

我使用的是方法是将所有数据一次性读入内存中，避免频繁进行磁盘IO，这样集中把所有数据读出来的时间要比一边训练一边读要快的多（使用较小的模型一般数据量不大，全部读入内存应该没什么问题，如果数据量较大呢？这时候用的模型一般也会较大，训练的时间占据主导，这时候就基本不会出现gpu等待数据的情况了）。以猫狗大战这个任务来说，自定义的Dataset如下，关键代码后用!!!..表示：

class CatDogDataset(Dataset):
    def __init__(self, root_dir, transform=None, test=False):
        self.root_dir = root_dir
        self.transform = transform
        self.image_paths = []
        self.image_data = []		# !!!!!!!!!!!!!!!!!!! 
        self.labels = []
        self.test = test

        for filename in os.listdir(root_dir):
            if filename.endswith('.jpg'):
                image_path = os.path.join(root_dir, filename)
                image = Image.open(image_path).convert('RGB')  # 转换为RGB格式
                if self.transform:
                    image = self.transform(image)

                self.image_paths.append(image_path)		
                self.image_data.append(image)		# !!!!!!!!!!!!!!!!!!!!	将所有图片读到内存进来
                if not test:
                    if 'cat' in filename:
                        self.labels.append(0)  # cat 类别标记为 0
                    elif 'dog' in filename:
                        self.labels.append(1)  # dog 类别标记为 1

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

    def __getitem__(self, idx):
        if self.test:
            return self.image_data[idx], self.image_paths[idx]  # 测试集返回图像及其路径
        else:
            return self.image_data[idx], self.labels[idx]