machine-learning - 如何在一次操作中将 torchvision.datasets 创建的数据集放入 GPU 中？

Question

我正在处理 CIFAR10，并使用 torchvision.datasets 来创建它。我需要 GPU 来加速计算，但我找不到一种方法将整个数据集一次性放入 GPU 中。我的模型需要使用小批量，单独处理每个批量确实非常耗时。

我尝试过将每个小批量分别放入 GPU，但这似乎非常耗时。

Answer 1

TL;DR

一次性移动整个数据集不会节省时间。

<小时/>

即使您有 GPU 内存来处理整个数据集，我认为您也不一定想要这样做(当然，按照今天的标准，CIFAR10 很小)。

我尝试了各种批量大小，并按如下方式定时传输到 GPU:

num_workers = 1 # Set this as needed

def time_gpu_cast(batch_size=1):
    start_time = time()
    for x, y in DataLoader(dataset, batch_size, num_workers=num_workers):
        x.cuda(); y.cuda()
    return time() - start_time

# Try various batch sizes
cast_times = [(2 ** bs, time_gpu_cast(2 ** bs)) for bs in range(15)]
# Try the entire dataset like you want to do
cast_times.append((len(dataset), time_gpu_cast(len(dataset))))

plot(*zip(*cast_times)) # Plot the time taken

对于num_workers = 1，这就是我得到的:

如果我们尝试并行加载(num_workers = 8)，它会变得更加清晰: