nvidia dali读数据测试程序

「爱情、让人受尽委屈。」 2022-12-13 11:27 74阅读 0赞

'''
    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)
            return [output, self.labels]
    
    
    class HybridValPipe(Pipeline):
        def __init__(self, batch_size, num_threads, device_id, data_dir, crop, size, local_rank=0, world_size=1):
            super(HybridValPipe, self).__init__(batch_size, num_threads, device_id, seed=12 + device_id)
            self.input = ops.FileReader(file_root=data_dir, shard_id=local_rank, num_shards=world_size,
                                        random_shuffle=False)
            self.decode = ops.ImageDecoder(device="mixed", output_type=types.RGB)
            self.res = ops.Resize(device="gpu", resize_shorter=size, interp_type=types.INTERP_TRIANGULAR)
            self.cmnp = ops.CropMirrorNormalize(device="gpu",
                                                output_dtype=types.FLOAT,
                                                output_layout=types.NCHW,
                                                crop=(crop, crop),
                                                image_type=types.RGB,
                                                mean=[0.485 * 255, 0.456 * 255, 0.406 * 255],
                                                std=[0.229 * 255, 0.224 * 255, 0.225 * 255])
    
        def define_graph(self):
            self.jpegs, self.labels = self.input(name="Reader")
            images = self.decode(self.jpegs)
            images = self.res(images)
            output = self.cmnp(images)
            return [output, self.labels]
    
    
    def get_imagenet_iter_dali(type, image_dir, batch_size, num_threads, device_id, num_gpus, crop, val_size=256,
                               world_size=1,
                               local_rank=0):
        if type == 'train':
            pip_train = HybridTrainPipe(batch_size=batch_size, num_threads=num_threads, device_id=local_rank,
                                        data_dir=image_dir + '/train',
                                        crop=crop, world_size=world_size, local_rank=local_rank)
            pip_train.build()
            dali_iter_train = DALIClassificationIterator(pip_train, size=pip_train.epoch_size("Reader") // world_size, auto_reset=True)
            return dali_iter_train
        elif type == 'val':
            pip_val = HybridValPipe(batch_size=batch_size, num_threads=num_threads, device_id=local_rank,
                                    data_dir=image_dir + '/val',
                                    crop=crop, size=val_size, world_size=world_size, local_rank=local_rank)
            pip_val.build()
            dali_iter_val = DALIClassificationIterator(pip_val, size=pip_val.epoch_size("Reader") // world_size, auto_reset=True)
            return dali_iter_val
    
    def getArgs():
        args = argparse.ArgumentParser()
        args.add_argument('-d', '--datadir', help='img folder containing train & val data',
                          default='data')  #数据目录，包含train，val两个目录，
        args.add_argument('-b', '--batch_size', help='batch size',type=int,
                          default='16')
        args.add_argument('-t', '--read_thread', help='num thread for reading data',type=int,
                          default='1') #加载数据的线程数
        return vars(args.parse_args())
    
    if __name__ == '__main__':
        args = getArgs()
        data_dir = args['datadir'] #数据目录
    
        batch_size = args['batch_size'] #bz
        num_workers = args['read_thread'] #读线程数
    
    
        train_loader = get_imagenet_iter_dali(type='train', image_dir=data_dir, batch_size=batch_size,
                                              num_threads=num_workers, crop=224, device_id=0, num_gpus=1)
        val_loader = get_imagenet_iter_dali(type='val', image_dir=data_dir, batch_size=batch_size,
                                              num_threads=num_workers, crop=224, device_id=0, num_gpus=1)
        num_epochs = 1
    
    
        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                dataloader = train_loader
            else:
                dataloader = val_loader
    
                # Iterate over data.
            for i, data in enumerate(dataloader):
                inputs = data[0]["data"].cuda(non_blocking=True)
                labels = data[0]["label"].squeeze().long().cuda(non_blocking=True)

ps：

*  auto\_reset参数置为True来实现多个epoch时数据的获取，否则每个epoch需要单独获取loader数据,如[这里][Link 1]。DALIClassificationIterator参数的设置详见[这里][Link 2]。
 *  文件夹中onehot的顺序根据[字母排序获得][Link 3]，即ls当前目录的结果。

FileReader supports a flat directory structure. The file_root directory must contain directories with data files. To obtain the labels, FileReader sorts directories in file_root in alphabetical order and takes an index in this order as a class label.

*  dali对部分数据无法处理，异常数据表现汇总如下

cannot identify image file

Corrupt EXIF data

libpng error: PNG input buffer is incomplete

*  可通过如下代码获取有问题图片

import os
    import imghdr
    import argparse
    from PIL import Image
    import warnings
    warnings.filterwarnings("error", category=UserWarning)
    
    try:
        # if True:
        img = Image.open(img_path)
        # img.verify()
        # print( imghdr.what(img_path), img.format)
        assert imghdr.what(img_path) != 'gif'
    
        assert imghdr.what(img_path) in ['jpg', 'jpeg', 'JPEG', 'png', 'bmp']
    
        if img.format in ['JPEG']:
            with open(img_path, 'rb') as f:
                check_chars = f.read()[-2:]
            # assert check_chars == b'\xff\xd9' #判断JPEG图片
    except Exception as e:
        # else:
        print ('{} error: {}'.format(img_path, str(e)))

参考文献：

1.  [DataLoaders\_DALI][DataLoaders_DALI]
2.  [NVIDIA DALI: Speeding up PyTorch][NVIDIA DALI_ Speeding up PyTorch]
3.  [nvidia.dali.ops.FileReader][Link 3]
4.  [Image Processing][]
5.  [nvidia-dali GPU加速预处理][nvidia-dali GPU]
6.  [PyTorch Plugin API reference ][Link 2]

[Link 1]: https://github.com/yaysummeriscoming/DALI_pytorch_demo/blob/master/main.py
[Link 2]: https://docs.nvidia.com/deeplearning/dali/user-guide/docs/plugins/pytorch_plugin_api.html
[Link 3]: https://docs.nvidia.com/deeplearning/dali/user-guide/docs/supported_ops.html#nvidia.dali.ops.FileReader
[DataLoaders_DALI]: https://github.com/tanglang96/DataLoaders_DALI
[NVIDIA DALI_ Speeding up PyTorch]: https://towardsdatascience.com/nvidia-dali-speeding-up-pytorch-876c80182440
[Image Processing]: https://docs.nvidia.com/deeplearning/dali/user-guide/docs/examples/image_processing/index.html
[nvidia-dali GPU]: https://www.twblogs.net/a/5d5f2731bd9eee5327fdec9a/?lang=zh-cn