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

时间:2020-10-07 05:58:58

标签: pytorch gpu

我相信使用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时,tqdm显示的“ it / s”成比例增加:

Plot of inference time vs batch size

我相信,只要不使用所有内存,GPU就可以同时处理整个张量。也许我对GPU如何并行处理数据一无所知,所以在此我将不胜感激。

我正在使用NVIDIA GeForce 2080 Ti,pytorch 1.6.0和CUDA 10.2。

1 个答案:

答案 0 :(得分:-1)

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