from nvidia.dali.pipeline import Pipeline from nvidia.dali import ops from nvidia.dali import types from nvidia.dali.plugin.pytorch import DALIClassificationIterator import numpy as np import torch from torch import nn class HybridTrainPipeline(Pipeline): def __init__(self, batch_size, file_root, num_threads, device_id, num_shards, shard_id): super(HybridTrainPipeline, self).__init__(batch_size, num_threads, device_id) device_type = {0:"cpu"} if num_shards == 0: self.input = ops.FileReader(file_root = file_root) else: self.input = ops.FileReader(file_root = file_root, num_shards = num_shards, shard_id = shard_id) # ##### 可自由更改 ################################### self.decode = ops.ImageDecoder(device = device_type.get(num_shards, "mixed"), output_type = types.RGB) self.res = ops.RandomResizedCrop(device=device_type.get(num_shards, "gpu"), size = 224) self.cmnp = ops.CropMirrorNormalize(device=device_type.get(num_shards, "gpu"), dtype = types.FLOAT, # output_dtype=types.FLOAT, output_layout=types.NCHW, 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]) # #################################################### def define_graph(self, ): jpegs, labels = self.input(name="Reader") images = self.decode(jpegs) images = self.res(images) images = self.cmnp(images) return images, labels class HybridValPipeline(Pipeline): def __init__(self, batch_size, file_root, num_threads, device_id, num_shards, shard_id): super(HybridValPipeline, self).__init__(batch_size, num_threads, device_id) device_type = {0:"cpu"} if num_shards == 0: self.input = ops.FileReader(file_root = file_root) else: self.input = ops.FileReader(file_root = file_root, num_shards = num_shards, shard_id = shard_id) # ##### 可自由更改 ################################### self.decode = ops.ImageDecoder(device = device_type.get(num_shards, "mixed"), output_type = types.RGB) self.res = ops.RandomResizedCrop(device=device_type.get(num_shards, "gpu"), size = 224) self.cmnp = ops.CropMirrorNormalize(device=device_type.get(num_shards, "gpu"), dtype = types.FLOAT, # output_dtype=types.FLOAT, output_layout=types.NCHW, 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]) # #################################################### def define_graph(self, ): jpegs, labels = self.input(name="Reader") images = self.decode(jpegs) images = self.res(images) images = self.cmnp(images) return images, labels class TorchWrapper: """ 将多个pipeline封装为一个iterator parameters: num_shards : int 显卡并行数 data_loader : dali.pipeline.Pipeline类型 经过pipeline处理的数据结果 iter_mode : str recursion, iter 指定多个pipeline合并的方式，默认recursion """ def __init__(self, num_shards, data_loader, iter_mode = "recursion"): self.index = 0 self.count = 0 self.num_shards = num_shards self.data_loader = data_loader self.iter_mode = iter_mode if self.iter_mode not in {"recursion", "iter"}: raise Exception("iter_mode should be either 'recursion' or 'iter'") def __iter__(self,): return self def __len__(self, ): # 返回样本总量，而非batch_num if num_shards == 0: return self.data_loader.size else: return len(self.data_loader)*self.data_loader[0].size def __next__(self, ): if num_shards == 0: # 不使用GPU data = next(self.data_loader) return data[0]["data"], data[0]["label"].view(-1).long() else: # 使用一块或多块GPU if self.iter_mode == "recursion": return self._get_next_recursion() elif self.iter_mode == "iter": return self._get_next_iter(self.data_loader[0]) def _get_next_iter(self, data_loader): if self.count == data_loader.size: self.index+=1 data_loader = self.data_loader[self.index] self.count+=1 data = next(data_loader) return data[0]["data"], data[0]["label"].view(-1).long() def _get_next_recursion(self, ): self.index = self.count%self.num_shards self.count+=1 data_loader = self.data_loader[self.index] data = next(data_loader) return data[0]["data"], data[0]["label"].view(-1).long() def get_iter_dali_cuda(batch_size=256, train_file_root="", val_file_root="", num_threads=4, device_id=[-1], num_shards=0, shard_id=[-1]): """ 获取可用于pytorch训练的数据迭代器 数据的读取和处理部分可以使用多张GPU来完成 1、创建dali pipeline 2、封装为适用于pytorch的数据迭代器 3、将多卡的各个pipeline封装在一起 4、数据输出在cpu端，在cuda中 数据需要保证如下形式： images |-file_list.txt |-images/dog |-dog_4.jpg |-dog_5.jpg |-dog_9.jpg |-dog_6.jpg |-dog_3.jpg |-images/kitten |-cat_10.jpg |-cat_5.jpg |-cat_9.jpg |-cat_8.jpg |-cat_1.jpg parameters: batch_size : int 每批数据的量 file_root : str 数据的路径 num_threads : int 读取数据的CPU线程数 device_id : list of int GPU的物理编号 shard_id : list of int GPU的虚拟编号 num_shard : int methods: get_train_pipeline(shard_id, device_id) : 创建dali的pipeline，用以读取并处理训练数据 get_val_pipeline(shard_id, device_id) : 创建dali的pipeline，用以读取并处理验证数据 get_dali_iter_for_torch(piplines, data_num) : 封装成可用于pytorch的数据迭代器 get_data_size(pipeline) : 计算每个pipeline实际输出的数据总量，数据总量是文件中的数据量，实际输出是去掉了不满一个批次大小的数据 例： # 分别从TRAIN_PATH和VAL_PATH读取训练和验证数据，batch_size选择256，启动4个线程来读取数据，用2块GPU处理数据，分别是第0号和第4号GPU # 程序默认使用所有显卡，和4线程 # 如果使用单张GPU，请设置num_shards = 1, shard_id = [0], device_id保持一个列表形式 # 如果不使用GPU，请使用get_iter_dali_cpu() train_data_iter, val_data_iter = get_iter_dali(batch_size=256, train_file_root=TRAIN_PATH, val_file_root=Val_PATH, num_threads=4, device_id=[0,4], num_shards=2, shard_id=[0,1]) # 在torch中训练 torch_model = TorchModel(para) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(torch_model.parameters()) for epoch in range(epoches): for step, x,y in enumerate(train_data_iter): # 数据 : x # 标签 : y x = x.to("cuda:0") y = y.to("cuda:0") output = my_model(x) optimizer.zero_grad() loss = criterion(output, y) loss.backward() optimizer.step() ... ... """ def get_train_pipeline(shard_id, device_id): pipeline = HybridTrainPipeline(batch_size = batch_size, file_root = train_file_root, num_threads = num_threads, num_shards = num_shards, shard_id = shard_id, device_id = device_id) return pipeline def get_val_pipeline(shard_id, device_id): pipeline = HybridValPipeline(batch_size = batch_size, file_root = val_file_root, num_threads = num_threads, num_shards = num_shards, shard_id = shard_id, device_id = device_id) return pipeline pipeline_for_train = [get_train_pipeline(shard_id = shard_id_index, device_id = device_id_index) \ for shard_id_index, device_id_index in zip(shard_id, device_id)] pipeline_for_val = [get_val_pipeline(shard_id = shard_id_index, device_id = device_id_index) \ for shard_id_index, device_id_index in zip(shard_id, device_id)] [pipeline.build() for pipeline in pipeline_for_train] [pipeline.build() for pipeline in pipeline_for_val] def get_data_size(pipeline): data_num = pipeline.epoch_size()["Reader"] batch_size = pipeline.batch_size return data_num//batch_size*batch_size data_num_train = get_data_size(pipeline_for_train[0]) data_num_val = get_data_size(pipeline_for_val[0]) def get_dali_iter_for_torch(pipelines, data_num): return [DALIClassificationIterator(pipelines=pipeline, last_batch_policy="drop",size = data_num) for pipeline in pipelines] data_loader_train = get_dali_iter_for_torch(pipeline_for_train, data_num_train) data_loader_val = get_dali_iter_for_torch(pipeline_for_val, data_num_val) train_data_iter = TorchWrapper(num_shards, data_loader_train) val_data_iter = TorchWrapper(num_shards, data_loader_val) return train_data_iter, val_data_iter def get_iter_dali_cpu(batch_size=256, train_file_root="", val_file_root="", num_threads=4): pipeline_train = HybridTrainPipeline(batch_size = batch_size, file_root = train_file_root, num_threads = num_threads, num_shards = 0, shard_id = -1, device_id = 0) pipeline_val = HybridTrainPipeline(batch_size = batch_size, file_root = val_file_root, num_threads = num_threads, num_shards = 0, shard_id = -1, device_id = 0) pipeline_train.build() pipeline_val.build() def get_data_size(pipeline): data_num = pipeline.epoch_size()["Reader"] batch_size = pipeline.batch_size return data_num//batch_size*batch_size data_num_train = get_data_size(pipeline_train) data_num_val = get_data_size(pipeline_val) data_loader_train = DALIClassificationIterator(pipelines=pipeline_train, last_batch_policy="drop",size = data_num_train) data_loader_val = DALIClassificationIterator(pipelines=pipeline_val, last_batch_policy="drop",size = data_num_val) train_data_iter = TorchWrapper(0,data_loader_train) val_data_iter = TorchWrapper(0,data_loader_val) return train_data_iter, val_data_iter if __name__ == "__main__": PATH = "./imagenet" TRAIN_PATH = "./imagenet/train" VALID_PATH = "./imagenet/val" train_data_iter_cuda, val_data_iter_cuda = get_iter_dali_cuda(batch_size=256, train_file_root=TRAIN_PATH, val_file_root=TRAIN_PATH, num_threads=4, device_id=[0,4], num_shards=2, shard_id=[0,1]) train_data_iter_cpu, val_data_iter_cpu = get_iter_dali_cpu(batch_size=256, train_file_root=TRAIN_PATH, val_file_root=TRAIN_PATH, num_threads=4)