PyTorch提速四倍！提高DALI利用率，創建基於CPU的Pipeline
- 2020 年 2 月 21 日
- 筆記
大數據文摘出品
來源：medium
編譯：趙吉克
在過去的幾年裡，深度學習硬體方面取得了巨大的進步，Nvidia的最新產品Tesla V100和Geforce RTX系列包含專用的張量核，用於加速神經網路中常用的操作。
特別值得一提的是，V100有足夠的能力以每秒數千張圖的速度訓練神經網路，這使得基於ImageNet數據集小模型在單GPU上訓練只需幾小時，與2012年在ImageNet上訓練AlexNet模型所花費的5天時間相比簡直是天壤之別!
然而，強大的GPU使數據預處理管道不堪重負。為了解決這個問題，Tensorflow發布了一個新的數據載入器：tf.data.Dataset，用C++編寫，並使用基於圖的方法將多個預處理操作鏈接在一起。
另一方面，PyTorch使用在PIL庫上用Python編寫的數據載入器，既方便優靈活，但在速度上有欠缺（儘管PIL-SIMD庫確實稍微改善了這種情況）。
進入NVIDIA數據載入器（DALI）：旨在消除數據預處理瓶頸，允許訓練和推理全速運行。DALI主要用於在GPU上的預處理，但是大多數操作也在CPU上有快速實現。本文主要關注PyTorch，但是DALI也支援Tensorflow、MXNet和TensorRT，尤其是TensorRT有高度支援。它允許訓練和推理步驟使用完全相同的預處理程式碼。需注意，不同的框架（如Tensorflow和PyTorch）通常在數據載入器之間有很小的差異，這可能會影響準確性。
本文是Medium上一位部落客展示了一些技術來提高DALI的使用率並創建了一個完全基於CPU的管道。這些技術用於保持長期的記憶體穩定，並且與DALI包提供的CPU和GPU管道相比，可以增加50%的批處理大小。
DALI長期記憶體使用
第一個問題是，RAM的使用隨著訓練時間的增加而增加，這會導致OOM錯誤(即使是在擁有78GB RAM的VM上)，並且尚未修正。
唯一解決方案是重新import DALI並每次重構訓練和驗證通道：
del self.train_loader, self.val_loader, self.train_pipe, self.val_pipe torch.cuda.synchronize() torch.cuda.empty_cache() gc.collect() importlib.reload(dali) from dali import HybridTrainPipe, HybridValPipe, DaliIteratorCPU, DaliIteratorGPU <rebuild DALI pipeline>
請注意，使用這種方法，DALI仍然需要大量RAM才能獲得最好的結果。考慮到如今RAM的價格，這並不是什麼大問題。從下表可以看出，DALI的最大批大小可能比TorchVision低50%：
接下來的部分涉及降低GPU佔用率的方法。
構建一個完全基於CPU的Pipeline
讓我們首先看看示例CPU管道。當不考慮峰值吞吐量時，基於CPU的管道非常有用。CPU訓練管道只在CPU上執行解碼和調整大小的操作，而CropMirrorNormalize操作則在GPU上運行。由於僅僅是傳輸輸出到GPU與DALI就使用了大量的GPU記憶體，為了避免這種情況，我們修改了示例CPU管道，使其完全運行在CPU上：
class HybridTrainPipe(Pipeline): def __init__(self, batch_size, num_threads, device_id, data_dir, crop, mean, std, local_rank=0, world_size=1, dali_cpu=False, shuffle=True, fp16=False, min_crop_size=0.08): # As we're recreating the Pipeline at every epoch, the seed must be -1 (random seed) super(HybridTrainPipe, self).__init__(batch_size, num_threads, device_id, seed=-1) # Enabling read_ahead slowed down processing ~40% self.input = ops.FileReader(file_root=data_dir, shard_id=local_rank, num_shards=world_size, random_shuffle=shuffle) # Let user decide which pipeline works best with the chosen model if dali_cpu: decode_device = "cpu" self.dali_device = "cpu" self.flip = ops.Flip(device=self.dali_device) else: decode_device = "mixed" self.dali_device = "gpu" output_dtype = types.FLOAT if self.dali_device == "gpu" and fp16: output_dtype = types.FLOAT16 self.cmn = ops.CropMirrorNormalize(device="gpu", output_dtype=output_dtype, output_layout=types.NCHW, crop=(crop, crop), image_type=types.RGB, mean=mean, std=std,) # To be able to handle all images from full-sized ImageNet, this padding sets the size of the internal nvJPEG buffers without additional reallocations device_memory_padding = 211025920 if decode_device == 'mixed' else 0 host_memory_padding = 140544512 if decode_device == 'mixed' else 0 self.decode = ops.ImageDecoderRandomCrop(device=decode_device, output_type=types.RGB, device_memory_padding=device_memory_padding, host_memory_padding=host_memory_padding, random_aspect_ratio=[0.8, 1.25], random_area=[min_crop_size, 1.0], num_attempts=100) # Resize as desired. To match torchvision data loader, use triangular interpolation. self.res = ops.Resize(device=self.dali_device, resize_x=crop, resize_y=crop, interp_type=types.INTERP_TRIANGULAR) self.coin = ops.CoinFlip(probability=0.5) print('DALI "{0}" variant'.format(self.dali_device)) def define_graph(self): rng = self.coin() self.jpegs, self.labels = self.input(name="Reader") # Combined decode & random crop images = self.decode(self.jpegs) # Resize as desired images = self.res(images) if self.dali_device == "gpu": output = self.cmn(images, mirror=rng) else: # CPU backend uses torch to apply mean & std output = self.flip(images, horizontal=rng) self.labels = self.labels.gpu() return [output, self.labels]DALI管道現在在CPU上輸出一個8位張量。我們需要使用PyTorch來完成CPU-> GPU的傳輸、浮點數的轉換和歸一化。這最後兩個操作是在GPU上完成的，快速並且減少了CPU -> GPU記憶體頻寬需求。在轉移到GPU之前，嘗試過固定這個張量，但是沒有得到任何性能提升，把它和一個預存器放在一起：
def _preproc_worker(dali_iterator, cuda_stream, fp16, mean, std, output_queue, proc_next_input, done_event, pin_memory): """ Worker function to parse DALI output & apply final preprocessing steps """ while not done_event.is_set(): # Wait until main thread signals to proc_next_input -- normally class="ql-long-20347411" style="line-height: 1.7;margin-bottom: 0pt;margin-top: 0pt;font-size: 11pt;color: #494949;"> proc_next_input.wait() proc_next_input.clear() if done_event.is_set(): print('Shutting down preproc thread') break try: data = next(dali_iterator) # Decode the data output input_orig = data[0]['data'] target = data[0]['label'].squeeze().long() # DALI should already output target style="line-height: 1.7;margin-bottom: 0pt;margin-top: 0pt;font-size: 11pt;color: #494949;"> # Copy to GPU and apply final processing in separate CUDA stream with torch.cuda.stream(cuda_stream): input = input_orig if pin_memory: input = input.pin_memory() del input_orig # Save memory input = input.cuda(non_blocking=True) input = input.permute(0, 3, 1, 2) # Input tensor is kept as 8-bit integer for transfer to GPU, to save bandwidth if fp16: input = input.half() else: input = input.float() input = input.sub_(mean).div_(std) # Put the result class="ql-long-20347411" style="line-height: 1.7;margin-bottom: 0pt;margin-top: 0pt;font-size: 11pt;color: #494949;"> output_queue.put((input, target)) except StopIteration: print('Resetting DALI loader') dali_iterator.reset() output_queue.put(None) class DaliIteratorCPU(DaliIterator): """ Wrapper class to decode the DALI iterator output & provide iterator that functions in the same way as TorchVision. Note that permutation to channels first, converting from 8-bit integer to float & normalization are all performed style="line-height: 1.7;margin-bottom: 0pt;margin-top: 0pt;font-size: 11pt;color: #494949;"> pipelines (Pipeline): DALI pipelines size (int): Number of examples in set fp16 (bool): Use fp16 as output format, f32 otherwise mean (tuple): Image mean value for each channel std (tuple): Image standard deviation value for each channel pin_memory (bool): Transfer input tensor to pinned memory, before moving to GPU """ def __init__(self, fp16=False, mean=(0., 0., 0.), std=(1., 1., 1.), pin_memory=True, **kwargs): super().__init__(**kwargs) print('Using DALI CPU iterator') self.stream = torch.cuda.Stream() self.fp16 = fp16 self.mean = torch.tensor(mean).cuda().view(1, 3, 1, 1) self.std = torch.tensor(std).cuda().view(1, 3, 1, 1) self.pin_memory = pin_memory if self.fp16: self.mean = self.mean.half() self.std = self.std.half() self.proc_next_input = Event() self.done_event = Event() self.output_queue = queue.Queue(maxsize=5) self.preproc_thread = threading.Thread( target=_preproc_worker, kwargs={'dali_iterator': self._dali_iterator, 'cuda_stream': self.stream, 'fp16': self.fp16, 'mean': self.mean, 'std': self.std, 'proc_next_input': self.proc_next_input, 'done_event': self.done_event, 'output_queue': self.output_queue, 'pin_memory': self.pin_memory}) self.preproc_thread.daemon = True self.preproc_thread.start() self.proc_next_input.set() def __next__(self): torch.cuda.current_stream().wait_stream(self.stream) data = self.output_queue.get() self.proc_next_input.set() if data is None: raise StopIteration return data def __del__(self): self.done_event.set() self.proc_next_input.set() torch.cuda.current_stream().wait_stream(self.stream) self.preproc_thread.join()基於GPU的Pipeline
測試中，在類似最大批處理大小下，上述CPU管道的速度大約是TorchVision數據載入器的兩倍。CPU管道可以很好地與像ResNet50這樣的大型模型一起工作；然而，當使用像AlexNet或ResNet18這樣的小模型時，CPU更好。GPU管道的問題是最大批處理大小減少了近50%，限制了吞吐量。
一種顯著減少GPU記憶體使用的方法是將驗證管道與GPU隔離直到最後再調用。這很容易做到，因為我們已經重新導入DALI，並在每個epoch中重新創建數據載入器。
更多小提示
在驗證時，將數據集均分的批處理大小效果最好，這避免了在驗證數據集結束時還需要進行不完整的批處理。
與Tensorflow和PyTorch數據載入器類似，TorchVision和DALI管道不會產生相同的輸出—您將看到驗證精度略有不同。我發現這是由於不同的JPEG影像解碼器。另一方面，DALI支援TensorRT，允許使用完全相同的預處理來進行訓練和推理。
對於峰值吞吐量，嘗試將數據載入器的數量設置為number_of_virtual_CPU核心，2個虛擬核對應1個物理核。
如果你想要絕對最好的性能並且不需要有類似TorchVision的輸出，嘗試關閉DALI三角形插值。
不要忘記磁碟IO。確保您有足夠的記憶體來快取數據集和/或一個非常快的SSD。DALI讀取高達400Mb/s !
合併
為了方便地集成這些修改，我創建了一個data loader類，其中包含這裡描述的所有修改，包括DALI和TorchVision後端。使用很簡單。實例化數據載入程式：
dataset = Dataset(data_dir, batch_size, val_batch_size workers, use_dali, dali_cpu, fp16)
然後得到訓練和驗證數據裝載器:
train_loader = dataset.get_train_loader() val_loader = dataset.get_val_loader()
在每個訓練周期結束時重置數據載入器:
dataset.reset()
或者，驗證管道可以在模型驗證之前在GPU上重新創建:
dataset.prep_for_val()
基準
以下是使用ResNet18的最大批量大小：
因此，通過應用這些修改，DALI可以在CPU和GPU模式下使用的最大批處理大小增加了約50%!
這裡是一些使用Shufflenet V2 0.5和批量大小512的吞吐量圖：
這裡是一些使用DALI GPU管道訓練各種網路，包括在TorchVision：
所有測試都在GoogleCloud V100實例上運行，該實例有12個vcpu(6個物理核心)，78GB RAM，並使用Apex FP16培訓。要重現這些結果，請使用以下參數：
— fp16 — batch-size 512 — workers 10 — arch 「shufflenet_v2_x0_5 or resnet18」 — prof — use-dali
所以，DALI使得單核特斯拉V100可以達到接近4000張/秒的影像處理速度！這達到了Nvidia DGX-1的一半多一點（它有8個V100 gpu），儘管我們使用了小模型。對我來說，能夠在幾個小時內在一個GPU上運行ImageNet是生產力進步。
本文中提供的程式碼如下：
https://github.com/yaysummeriscoming/DALI_pytorch_demo
相關報道：
https://towardsdatascience.com/nvidia-dali-speeding-up-pytorch-876c80182440
