nvidia dali读数据测试程序

最新推荐文章于 2025-06-08 09:01:31 发布

hzhj

最新推荐文章于 2025-06-08 09:01:31 发布

阅读量1.6k

点赞数

CC 4.0 BY-SA版权

分类专栏：图像处理文章标签： nvidia dali image process

本文链接：https://blog.csdn.net/hzhj2007/article/details/109009825

本文介绍了如何使用NVIDIA DALI在多个epoch中读取数据，重点强调了`auto_reset`参数的作用。内容涉及DALIClassificationIterator的配置，以及如何处理文件夹中按字母顺序排列的onehot数据。同时，文章列举了DALI在处理数据时可能出现的异常情况，如无法识别的图像文件、损坏的EXIF数据和不完整的PNG输入缓冲区，并提供了识别问题图片的代码。参考资源包括NVIDIA DALI的PyTorch插件API和图像处理指南。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

'''
python test_dali.py -d data/ -b 64 -t 4

-d 数据地址
-b batchsize
-t 读数据的线程数
'''
import argparse
import nvidia.dali.ops as ops
import nvidia.dali.types as types
from nvidia.dali.pipeline import Pipeline
from nvidia.dali.plugin.pytorch import DALIClassificationIterator, DALIGenericIterator

class HybridTrainPipe(Pipeline):
    def __init__(self, batch_size, num_threads, device_id, data_dir, crop, dali_cpu=False, local_rank=0, world_size=1):
        super(HybridTrainPipe, self).__init__(batch_size, num_threads, device_id, seed=12 + device_id)
        dali_device = "gpu"
        self.input = ops.FileReader(file_root=data_dir, shard_id=local_rank, num_shards=world_size, random_shuffle=True)
        self.decode = ops.ImageDecoder(device="mixed", output_type=types.RGB)
        self.res = ops.RandomResizedCrop(device="gpu", size=crop, random_area=[0.08, 1.25])
        self.cmnp = ops.CropMirrorNormalize(device="gpu",
                                            output_dtype=types.FLOAT,
                                            output_layout=types.NCHW,
                                            image_type=types.RGB,
                                            mean=[0.485 * 255, 0.456 * 255, 0.406 * 255],
                                            std=[0.229 * 255, 0.224 * 255, 0.225 * 255])
        self.coin = ops.CoinFlip(probability=0.5)
        print('DALI "{0}" variant'.format(dali_device))

    def define_graph(self):
        rng = self.coin()
        self.jpegs, self.labels = self.input(name="Reader")
        images = self.decode(self.jpegs)
        images = self.res(images)
        output = self.cmnp(images, mirror=rng)