在推理过程中增加批量大小

Question

我相信当使用GPU时，每个批次的推理时间与批次大小无关，但这个最小的例子告诉我这似乎不是真的：

import torch
from torch import nn
from tqdm import tqdm
BATCH_SIZE = 32
N_ITER = 10000
class NN(nn.Module):
def __init__(self):
super(NN, self).__init__()
self.layer = nn.Conv2d(3, 32, kernel_size=5, stride=1, padding=3, bias=False)
def forward(self, input):
out = self.layer(input)
return out

cnn = NN().cuda()
cnn.eval()
tensor = torch.rand(BATCH_SIZE, 3, 999, 999).cuda()
with torch.no_grad():
for _ in tqdm(range(N_ITER), mininterval=0.1):
out = cnn(tensor)

当增加BATCH_ SIZE时；它/s"；tqdm按比例增加：

推断时间与批量的关系图

我相信GPU可以同时处理整个张量，只要它不使用所有的内存。也许我不了解GPU是如何并行处理数据的，所以我很感激这里的一些见解。

我使用的是NVIDIA GeForce 2080 Ti、pytorch 1.6.0和CUDA 10.2。

Answer 1

你错了。GPU有很多核心，但这并不意味着它们可以同时处理所有数据。例如，RTX 2080Ti只有4352个核心。