如何将整个数据集加载到GPU

发布于 2025-02-12 08:25:18 字数 129 浏览 4 评论 0原文

我有1550张图像的数据集3x112x112。当我训练模型时，我会创建数据集WIA ImageFolder，然后使用DataLoader。由于每次从内存阅读，这需要花费太多时间。我有足够的GPU内存一次加载整个数据集。会更快。最好的方法是什么？

原文

分享到QQ

分享到微博

如果你对这篇内容有疑问，欢迎到本站社区发帖提问参与讨论，获取更多帮助，或者扫码二维码加入 Web 技术交流群。

发布评论

需要登录才能够评论，你可以免费注册一个本站的账号。

空宴 2025-02-19 08:25:18

您可以将图像作为数据集的属性存储，在初始化时将其放在GPU上，然后让__ getItem __直接从中返回图像。

import torch
from torch.utils.data import Dataset, DataLoader

class GPUDataset(Dataset):
    def __init__(self):
        self.len = 1550
        self.data = torch.randn(self.len, 3, 112, 112).cuda()  # ~ 223 MB

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len
    
dataset = GPUDataset()
loader = DataLoader(dataset, batch_size=4)

batch = next(iter(loader))
print(batch.device, batch.shape)    # cuda:0 torch.Size([4, 3, 112, 112])

这是针对小型数据集的快速修复。对于较大的数据集，请考虑更高级的解决方案，例如 lmdb 。

You can store the images as an attribute of the dataset, put it on the GPU at initialization, and let __getitem__ return images from this directly.

import torch
from torch.utils.data import Dataset, DataLoader

class GPUDataset(Dataset):
    def __init__(self):
        self.len = 1550
        self.data = torch.randn(self.len, 3, 112, 112).cuda()  # ~ 223 MB

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len
    
dataset = GPUDataset()
loader = DataLoader(dataset, batch_size=4)

batch = next(iter(loader))
print(batch.device, batch.shape)    # cuda:0 torch.Size([4, 3, 112, 112])

This is meant to be a quick fix for small datasets. For larger datasets, consider more advanced solutions such as lmdb.

回复收藏 0 原文

~没有更多了~