torch中dataloader加速

最新推荐文章于 2024-07-02 21:03:46 发布

Nicola-Zhang

最新推荐文章于 2024-07-02 21:03:46 发布

阅读量873

点赞数

分类专栏： torch

本文链接：https://blog.csdn.net/yangwangnndd/article/details/102134187

版权

本文分享了在使用Torch DataLoader进行数据加载时的实际测试经验，指出在处理ImageNet数据集时，加速效果并不明显，甚至可能导致代码复杂度增加和效率下降。

摘要由CSDN通过智能技术生成

注意：经过实际测试，在ImageNet数据集上并没有达到所说的提升效果，反而会导致代码逻辑混乱，徒增代码量！

class data_prefetcher():
    def __init__(self, loader):
        self.loader = iter(loader)
        self.stream = torch.cuda.Stream()
        self.mean = torch.tensor([0.485 * 255, 0.456 * 255, 0.406 * 255]).cuda().view(1,3,1,1)
        self.std = torch.tensor([0.229 * 255, 0.224 * 255, 0.225 * 255]).cuda().view(1,3,1,1)
        # With Amp, it isn't necessary to manually convert data to half.
        # if args.fp16:
        #     self.mean = self.mean.half()
        #     self.std = self.std.half()
        self.preload()

    def preload(self):
        try:
            self.next_input, self.next_target = next(self.loader)
        except StopIteration:
            self.next_input = None
            self.next_target = None
            return
        with torch.cuda.stream(self.stream):
            self.next_input &