时间戳

首页 > 解决方案 > 几次迭代后，Pytorch 速度变慢

python - 几次迭代后，Pytorch 速度变慢

问题描述

我正在尝试在 PyTorch 中实现一个模型。训练过程相当复杂，需要一些时间，但我注意到模型在前几批非常快，然后突然达到 500 左右。我猜是由于一些内存泄漏问题，好像 python并没有真正释放释放的巨大张量的内存。

起初我以为问题与存储梯度有关，但实际上即使torch.no_grad()出现了同样的问题。

这是一个复制问题的示例。（注意我不是在尝试训练这个特定的网络，但问题看起来是一样的）。为了使事情更简单，我没有使用渐变，而是在同一批次上进行迭代。

import torch
import torch.nn as nn
from torchvision.datasets import MNIST
import torchvision.transforms as T

dataset = MNIST(root='./MNIST', train=True, download=True,
                transform=T.Compose([T.ToTensor(), T.Lambda(lambda x: torch.flatten(x))]))
data_loader = torch.utils.data.DataLoader(dataset, batch_size=500)

X, _ = next(iter(data_loader))
X = X.to('cuda')

in_features = 28*28
out_features = 10
width= 15000

#defining huge network
NN = nn.Sequential(
          nn.Linear(in_features=28*28, out_features=width, bias=False),
          nn.ReLU(),
          nn.Linear(in_features=width, out_features=width, bias=False),
          nn.ReLU(),
          nn.Linear(in_features=width, out_features=width, bias=False),
          nn.ReLU(),
          nn.Linear(in_features=width, out_features=width, bias=False),
          nn.ReLU(),
          nn.Linear(in_features=width, out_features=width, bias=False),
          nn.ReLU(),
          nn.Linear(in_features=width, out_features=width, bias=False),
          nn.ReLU(),
          nn.Linear(in_features=width, out_features=width, bias=False),
          nn.ReLU(),
          nn.Linear(in_features=width, out_features=width, bias=False),
          nn.ReLU(),
          nn.Linear(in_features=width, out_features=width, bias=False),
          nn.ReLU(),
          nn.Linear(in_features=width, out_features=width, bias=False),
          nn.ReLU(),
          nn.Linear(in_features=width, out_features=out_features, bias=False),
        ).to('cuda')

import time

iterations=100

X = X.to('cuda')

with torch.no_grad():
  for idx in range(iterations):
    print(f'Iteration {idx+1}')
    start = time.time()
    Y = NN(X)
    print(f'Time: {time.time() - start}')

输出显示一切都非常快，直到第 50 次迭代，然后它突然变慢了。

Iteration 44
Time: 0.00035953521728515625
Iteration 45
Time: 0.00035309791564941406
Iteration 46
Time: 0.00035309791564941406
Iteration 47
Time: 0.048192501068115234
Iteration 48
Time: 0.1714644432067871
Iteration 49
Time: 0.16771984100341797
Iteration 50
Time: 0.1681973934173584
Iteration 51
Time: 0.16853046417236328
Iteration 52
Time: 0.16821908950805664

为什么会出现这样的减速？有没有可能以某种方式避免它？

标签： pythonfor-loopmemory-leakspytorchgpu

解决方案

查看此页面并向下滚动到“异步执行”。

基本上，您测量的是将操作排入 GPU 的时间，而不是实际执行操作所需的时间。这是因为 GPU 调用是异步的，如链接中所述。我复制了下面的相关部分：

默认情况下，GPU 操作是异步的。当您调用使用 GPU 的函数时，这些操作将被排入特定设备的队列，但不一定要等到稍后才会执行。这使我们能够并行执行更多计算，包括在 CPU 或其他 GPU 上的操作。

一般来说，异步计算的效果对调用者来说是不可见的，因为（1）每个设备按照它们排队的顺序执行操作，（2）PyTorch 在 CPU 和 GPU 之间或两个 GPU 之间复制数据时会自动执行必要的同步。因此，计算将像每个操作都同步执行一样进行。

您可以通过设置环境变量 CUDA_LAUNCH_BLOCKING=1 来强制同步计算。当 GPU 发生错误时，这会很方便。（对于异步执行，直到操作实际执行后才会报告此类错误，因此堆栈跟踪不会显示请求的位置。）

异步计算的结果是没有同步的时间测量是不准确的。要获得精确的测量值，应该在测量之前调用 torch.cuda.synchronize()，或者使用torch.cuda.Event如下方式记录时间：

start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
start_event.record()

# Run some things here

end_event.record()
torch.cuda.synchronize()  # Wait for the events to be recorded!
elapsed_time_ms = start_event.elapsed_time(end_event)

推荐阅读