在利用DL解決圖像問題時,影響訓練效率最大的有時候是GPU,有時候也可能是CPU和你的磁盤。
很多設計不當的任務,在訓練神經網絡的時候,大部分時間都是在從磁盤中讀取數據,而不是做 Backpropagation 。
這種癥狀的體現是使用 Nividia-smi 查看 GPU 使用率時,Memory-Usage 占用率很高,但是 GPU-Util 時常為 0% ,如下圖所示:

如何解決這種問題呢?
在 Nvidia 提出的分布式框架 Apex 里面,我們在源碼里面找到了一個簡單的解決方案:
https://github.com/NVIDIA/apex/blob/f5cd5ae937f168c763985f627bbf850648ea5f3f/examples/imagenet/main_amp.py#L256
class data_prefetcher():
def __init__(self, loader):
self.loader = iter(loader)
self.stream = torch.cuda.Stream()
self.mean = torch.tensor([0.485 * 255, 0.456 * 255, 0.406 * 255]).cuda().view(1,3,1,1)
self.std = torch.tensor([0.229 * 255, 0.224 * 255, 0.225 * 255]).cuda().view(1,3,1,1)
# With Amp, it isn't necessary to manually convert data to half.
# if args.fp16:
# self.mean = self.mean.half()
# self.std = self.std.half()
self.preload()
def preload(self):
try:
self.next_input, self.next_target = next(self.loader)
except StopIteration:
self.next_input = None
self.next_target = None
return
with torch.cuda.stream(self.stream):
self.next_input = self.next_input.cuda(non_blocking=True)
self.next_target = self.next_target.cuda(non_blocking=True)
# With Amp, it isn't necessary to manually convert data to half.
# if args.fp16:
# self.next_input = self.next_input.half()
# else:
self.next_input = self.next_input.float()
self.next_input = self.next_input.sub_(self.mean).div_(self.std)
我們能看到 Nvidia 是在讀取每次數據返回給網絡的時候,預讀取下一次迭代需要的數據,
那么對我們自己的訓練代碼只需要做下面的改造:
training_data_loader = DataLoader(
dataset=train_dataset,
num_workers=opts.threads,
batch_size=opts.batchSize,
pin_memory=True,
shuffle=True,
)
for iteration, batch in enumerate(training_data_loader, 1):
# 訓練代碼
#-------------升級后---------
data, label = prefetcher.next()
iteration = 0
while data is not None:
iteration += 1
# 訓練代碼
data, label = prefetcher.next()
這樣子我們的 Dataloader 就像打了雞血一樣提高了效率很多,如下圖:

當然,最好的解決方案還是從硬件上,把讀取速度慢的機械硬盤換成 NVME 固態吧~
補充:Pytorch設置多線程進行dataloader時影響GPU運行
使用PyTorch設置多線程(threads)進行數據讀取時,其實是假的多線程,他是開了N個子進程(PID是連續的)進行模擬多線程工作。
以載入cocodataset為例
DataLoader
dataloader = torch.utils.data.DataLoader(COCODataset(config["train_path"],
(config["img_w"], config["img_h"]),
is_training=True),
batch_size=config["batch_size"],
shuffle=True, num_workers=32, pin_memory=True)
numworkers就是指定多少線程的參數,原為32。
檢查GPU是否運行該程序
查看運行在gpu上的所有程序:
如果沒有返回,則該程序并沒有在GPU上運行
指定GPU運行
將num_workers改成0即可
以上為個人經驗,希望能給大家一個參考,也希望大家多多支持腳本之家。
您可能感興趣的文章:- pytorch鎖死在dataloader(訓練時卡死)
- pytorch Dataset,DataLoader產生自定義的訓練數據案例
- 解決Pytorch dataloader時報錯每個tensor維度不一樣的問題
- pytorch中DataLoader()過程中遇到的一些問題
- Pytorch dataloader在加載最后一個batch時卡死的解決
- pytorch DataLoader的num_workers參數與設置大小詳解
- pytorch 實現多個Dataloader同時訓練