|
- import os
-
- import nvidia.dali.ops as ops
- import nvidia.dali.types as types
- import torch.utils.data
- from nvidia.dali.pipeline import Pipeline
- from nvidia.dali.plugin.pytorch import DALIClassificationIterator
-
-
- class HybridTrainPipe(Pipeline):
- def __init__(self, batch_size, num_threads, device_id, data_dir, crop, seed=12, local_rank=0, world_size=1,
- spos_pre=False):
- super(HybridTrainPipe, self).__init__(batch_size, num_threads, device_id, seed=seed + device_id)
- color_space_type = types.BGR if spos_pre else types.RGB
- 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.BGR) # color_space_type
- self.res = ops.RandomResizedCrop(device="gpu", size=crop,
- interp_type=types.INTERP_LINEAR if spos_pre else types.INTERP_TRIANGULAR)
- self.twist = ops.ColorTwist(device="gpu")
- self.jitter_rng = ops.Uniform(range=[0.6, 1.4])
- # self.cmnp = ops.CropMirrorNormalize(device="gpu",
- # dtype = types.FLOAT, # output_dtype=types.FLOAT,
- # output_layout=types.NCHW,
- # # image_type=color_space_type, # 该功能被删掉了，在ImageDecoder中即可完成
- # mean=0. if spos_pre else [0.485 * 255, 0.456 * 255, 0.406 * 255],
- # std=1. if spos_pre else [0.229 * 255, 0.224 * 255, 0.225 * 255])
- self.cmnp = ops.CropMirrorNormalize(device="gpu", dtype = types.FLOAT, output_layout=types.NCHW,
- 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)
-
- def define_graph(self):
- rng = self.coin()
- self.jpegs, self.labels = self.input(name="Reader")
- images = self.decode(self.jpegs)
- images = self.res(images)
- images = self.twist(images, saturation=self.jitter_rng(),
- contrast=self.jitter_rng(), brightness=self.jitter_rng())
- output = self.cmnp(images, mirror=rng) # 临时删除，测试准确率为零是否是数据处理的原因
- return [output, self.labels] # output
-
-
- class HybridValPipe(Pipeline):
- def __init__(self, batch_size, num_threads, device_id, data_dir, crop, size, seed=12, local_rank=0, world_size=1,
- spos_pre=False, shuffle=False):
- super(HybridValPipe, self).__init__(batch_size, num_threads, device_id, seed=seed + device_id)
- color_space_type = types.BGR if spos_pre else types.RGB
- self.input = ops.FileReader(file_root=data_dir, shard_id=local_rank, num_shards=world_size,
- random_shuffle=shuffle)
- self.decode = ops.ImageDecoder(device="mixed", output_type=types.BGR)
- self.res = ops.Resize(device="gpu", resize_shorter=size,
- interp_type=types.INTERP_LINEAR if spos_pre else types.INTERP_TRIANGULAR)
- # self.cmnp = ops.CropMirrorNormalize(device="gpu",
- # dtype = types.FLOAT, # output_dtype=types.FLOAT,
- # output_layout=types.NCHW,
- # crop=(crop, crop),
- # # image_type=color_space_type,
- # mean=0. if spos_pre else [0.485 * 255, 0.456 * 255, 0.406 * 255],
- # std=1. if spos_pre else [0.229 * 255, 0.224 * 255, 0.225 * 255])
- self.cmnp = ops.CropMirrorNormalize(device="gpu",
- dtype = types.FLOAT, # output_dtype=types.FLOAT,
- output_layout=types.NCHW,
- crop=(crop, crop),
- # image_type=color_space_type,
- 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]
-
-
- class ClassificationWrapper:
- def __init__(self, loader, size):
- self.loader = loader
- self.size = size
-
- def __iter__(self):
- return self
-
- def __next__(self):
- data = next(self.loader)
- return data[0]["data"], data[0]["label"].view(-1).long().cuda(device="cuda:0", non_blocking=True) # .cuda(non_blocking=True)
-
- def __len__(self):
- return self.size
-
-
- def get_imagenet_iter_dali(split, image_dir, batch_size, num_threads, crop=224, val_size=256,
- spos_preprocessing=False, seed=12, shuffle=False, device_id=None):
- world_size, local_rank = 1, 0
- if device_id is None:
- device_id = torch.cuda.device_count() - 1 # use last gpu
- if split == "train":
- pipeline = HybridTrainPipe(batch_size=batch_size, num_threads=num_threads, device_id=device_id,
- data_dir=os.path.join(image_dir, "train"), seed=seed,
- crop=crop, world_size=world_size, local_rank=local_rank,
- spos_pre=spos_preprocessing)
- elif split == "val":
- pipeline = HybridValPipe(batch_size=batch_size, num_threads=num_threads, device_id=device_id,
- data_dir=os.path.join(image_dir, "val"), seed=seed,
- crop=crop, size=val_size, world_size=world_size, local_rank=local_rank,
- spos_pre=spos_preprocessing, shuffle=shuffle)
- else:
- raise AssertionError
- pipeline.build()
- num_samples = pipeline.epoch_size("Reader")
- # fill_last_batch的设置
- # 参考这里, valid和train设置为一样的策略
- # https://docs.nvidia.com/deeplearning/dali/user-guide/docs/plugins/pytorch_plugin_api.html?highlight=daliclassificationiterator#nvidia.dali.plugin.pytorch.DALIClassificationIterator
- last_batch_policy = ""
- last_batch_padded = True
- return ClassificationWrapper(
- DALIClassificationIterator(pipeline,
- # size=num_samples,
- last_batch_policy = last_batch_policy,
- # last_batch_padded = last_batch_padded,
- # fill_last_batch=split == "train", # 这个方法已经不建议使用了
- auto_reset=True),
- (num_samples + batch_size - 1) // batch_size)
|
