相关疑难解决方法(0)

我一直在尝试各种选项来加速 PyTorch 中的一些 for 循环逻辑。执行此操作的两个明显选项是使用numba或编写自定义 C++ 扩展。

作为一个例子，我从数字信号处理中选择了“可变长度延迟线”。使用简单的 Python for 循环可以简单但低效地编写此代码：

def delay_line(samples, delays):\n    """\n    :param samples: Float tensor of shape (N,)\n    :param delays: Int tensor of shape (N,)\n    \n    The goal is basically to mix each `samples[i]` with the delayed sample\n    specified by a per-sample `delays[i]`.\n    """\n    for i in range(len(samples)):\n        delay = int(delays[i].item())\n        index_delayed = i - delay\n        if index_delayed < 0:\n            index_delayed = 0\n\n        samples[i] = 0.5 * (samples[i] + samples[index_delayed])\n …

我试图找出GPU张量操作实际上是否比CPU更快.所以,我在下面编写了这个特殊的代码来连续实现CPU张量和GPU cuda张量的简单2D添加,以查看速度差异:

import torch
import time

###CPU
start_time = time.time()
a = torch.ones(4,4)
for _ in range(1000000):
    a += a
elapsed_time = time.time() - start_time

print('CPU time = ',elapsed_time)

###GPU
start_time = time.time()
b = torch.ones(4,4).cuda()
for _ in range(1000000):
    b += b
elapsed_time = time.time() - start_time

print('GPU time = ',elapsed_time)

令我惊讶的是,CPU时间为0.93秒,GPU时间高达63秒.我是否正确地进行了cuda张量操作,或者cuda张量的概念是否仅在非常复杂的操作中更快地运行,如在神经网络中？

注意:我的GPU是NVIDIA 940MX,torch.cuda.is_available()调用返回True.