Мне интересно, почему закрепление памяти в PyTorch делает вещи еще медленнее. Читая код torch.utils.data.dataloader, я обнаружил, что опция pin_memory=True для DataLoader просто вызывает .pin_memory() в каждом пакете, прежде чем вернуть их. Возвращенный тензор все еще находится на CPU, и после этого мне нужно вызвать .cuda(non_blocking=True) вручную. Таким образом, весь процесс будет
for x in some_iter:
yield x.pin_memory().cuda(non_blocking=True)
Я сравнил производительность этого с
for x in some_iter:
yield x.cuda()
Вот фактический код
a = torch.rand(1024, 655360)
%%time
for i in a:
i.pin_memory().cuda(non_blocking=True)
# CPU times: user 1.35 s, sys: 55.8 ms, total: 1.41 s
# Wall time: 396 ms
%%time
for i in a:
i.pin_memory().cuda()
# CPU times: user 1.6 s, sys: 12.2 ms, total: 1.62 s
# Wall time: 404 ms
%%time
for i in a:
i.cuda(non_blocking=True)
# CPU times: user 855 ms, sys: 3.87 ms, total: 859 ms
# Wall time: 274 ms
%%time
for i in a:
i.cuda()
# CPU times: user 314 ms, sys: 12 µs, total: 314 ms
# Wall time: 313 ms
В результате, не закрепляет память и использует меньше процессорного времени, и быстрее с точки зрения фактического времени . Не должно ли закрепление памяти сделать передачу данных асинхронной и, следовательно, быть быстрее? Если это не так, зачем нам пин-память?
PS. Я думал о возможности закрепления целых TensorDataset заранее (вместо того, чтобы закреплять партии каждый раз). Но это не может закрепить тензор, который больше, чем память GPU:
a = np.memmap('../dat/R/train.3,31,31B', '3,31,31B', 'r')
a.nbytes // 2**30
## 68
torch.from_numpy(a).pin_memory()
## ---------------------------------------------------------------------------
## RuntimeError Traceback (most recent call last)
## <ipython-input-36-d6f2d74da8e7> in <module>
## ----> 1 torch.from_numpy(a).pin_memory()
##
## RuntimeError: cuda runtime error (2) : out of memory at /tmp/pip-req-build-58y_cjjl/aten/src/THC/THCCachingHostAllocator.cpp:296
И если я действительно хочу закрепить небольшой тензор, почему бы мне не перенести непосредственно весь тензор в память GPU заранее?