使用 NumPy 和 PyTorch 在 GPU 上求解线性方程

Question

我正在尝试尽快求解大量线性方程。为了找出最快的方法，我在 CPU 和 GeForce 1080 GPU 上对NumPy和PyTorch进行了基准测试（使用Numba进行 NumPy）。结果真的让我很困惑。

这是我在 Python 3.8 中使用的代码：

import timeit

import torch
import numpy
from numba import njit


def solve_numpy_cpu(dim: int = 5):
    a = numpy.random.rand(dim, dim)
    b = numpy.random.rand(dim)

    for _ in range(1000):
        numpy.linalg.solve(a, b)


def solve_numpy_njit_a(dim: int = 5):
    njit(solve_numpy_cpu, dim=dim)


@njit
def solve_numpy_njit_b(dim: int = 5):
    a = numpy.random.rand(dim, dim)
    b = numpy.random.rand(dim)

    for _ in range(1000):
        numpy.linalg.solve(a, b)


def solve_torch_cpu(dim: int = 5):
    a = torch.rand(dim, dim)
    b = torch.rand(dim, 1)

    for _ in range(1000):
        torch.solve(b, a)


def solve_torch_gpu(dim: int = 5):
    torch.set_default_tensor_type("torch.cuda.FloatTensor")
    solve_torch_cpu(dim=dim)


def main():
    for f in (solve_numpy_cpu, solve_torch_cpu, solve_torch_gpu, solve_numpy_njit_a, solve_numpy_njit_b):
        time = timeit.timeit(f, number=1)
        print(f"{f.__name__:<20s}: {time:f}")


if __name__ == "__main__":
    main()

这些是结果：

solve_numpy_cpu     : 0.007275
solve_torch_cpu     : 0.012244
solve_torch_gpu     : 5.239126
solve_numpy_njit_a  : 0.000158
solve_numpy_njit_b  : 1.273660

最慢的是 CUDA 加速的 PyTorch。我验证了 PyTorch 正在使用我的 GPU

import torch
torch.cuda.is_available()
torch.cuda.get_device_name(0)

返回

True
'GeForce GTX 1080'

我可以理解，在 CPU 上，PyTorch 比 NumPy 慢。我不明白的是为什么 GPU 上的 PyTorch 慢得多。不是那么重要，但实际上更令人困惑的是，Numba 的njit装饰器会使性能降低几个数量级，直到您不再使用 @ 装饰器语法。

是我的设置吗？有时我会收到一条关于 Windows 页面/交换文件不够大的奇怪消息。如果我采用了一条完全晦涩的路径来在 GPU 上求解线性方程，我很乐意被引导到另一个方向。

---

编辑

因此，我专注于 Numba 并稍微改变了我的基准测试。正如 @max9111 所建议的，我重写了接收输入并生成输出的函数，因为最终，这就是任何人都想使用它们的目的。现在，我还对 Numba 加速函数执行第一次编译运行，因此后续计时更加公平。最后，我根据矩阵大小检查了性能并绘制了结果。

TL/DR：在矩阵大小达到 500x500 的情况下，Numba 加速对于numpy.linalg.solve.

这是代码：

import time
from typing import Tuple

import numpy
from matplotlib import pyplot
from numba import jit


@jit(nopython=True)
def solve_numpy_njit(a: numpy.ndarray, b: numpy.ndarray) -> numpy.ndarray:
    parameters = numpy.linalg.solve(a, b)
    return parameters


def solve_numpy(a: numpy.ndarray, b: numpy.ndarray) -> numpy.ndarray:
    parameters = numpy.linalg.solve(a, b)
    return parameters


def get_data(dim: int) -> Tuple[numpy.ndarray, numpy.ndarray]:
    a = numpy.random.random((dim, dim))
    b = numpy.random.random(dim)
    return a, b


def main():
    a, b = get_data(10)
    # compile numba function
    p = solve_numpy_njit(a, b)

    matrix_size = [(x + 1) * 10 for x in range(50)]
    non_accelerated = []
    accelerated = []
    results = non_accelerated, accelerated

    for j, each_matrix_size in enumerate(matrix_size):
        for m, f in enumerate((solve_numpy, solve_numpy_njit)):
            average_time = -1.
            for k in range(5):
                time_start = time.time()
                for i in range(100):
                    a, b = get_data(each_matrix_size)
                    p = f(a, b)
                d_t = time.time() - time_start
                print(f"{each_matrix_size:d} {f.__name__:<30s}: {d_t:f}")
                average_time = (average_time * k + d_t) / (k + 1)
            results[m].append(average_time)

    pyplot.plot(matrix_size, non_accelerated, label="not numba")
    pyplot.plot(matrix_size, accelerated, label="numba")
    pyplot.legend()
    pyplot.show()


if __name__ == "__main__":
    main()

这些是结果（针对矩阵边长的运行时间）：

---

编辑2

看到 Numba 对我的情况没有太大影响，我又开始对 PyTorch 进行基准测试。事实上，即使不使用 CUDA 设备，它也比 Numpy 快大约 4 倍。

这是我使用的代码：

import time
from typing import Tuple

import numpy
import torch
from matplotlib import pyplot


def solve_numpy(a: numpy.ndarray, b: numpy.ndarray) -> numpy.ndarray:
    parameters = numpy.linalg.solve(a, b)
    return parameters


def get_data(dim: int) -> Tuple[numpy.ndarray, numpy.ndarray]:
    a = numpy.random.random((dim, dim))
    b = numpy.random.random(dim)
    return a, b


def get_data_torch(dim: int) -> Tuple[torch.tensor, torch.tensor]:
    a = torch.rand(dim, dim)
    b = torch.rand(dim, 1)
    return a, b


def solve_torch(a: torch.tensor, b: torch.tensor) -> torch.tensor:
    parameters, _ = torch.solve(b, a)
    return parameters


def experiment_numpy(matrix_size: int, repetitions: int = 100):
    for i in range(repetitions):
        a, b = get_data(matrix_size)
        p = solve_numpy(a, b)


def experiment_pytorch(matrix_size: int, repetitions: int = 100):
    for i in range(repetitions):
        a, b = get_data_torch(matrix_size)
        p = solve_torch(a, b)


def main():
    matrix_size = [x for x in range(5, 505, 5)]
    experiments = experiment_numpy, experiment_pytorch
    results = tuple([] for _ in experiments)

    for i, each_experiment in enumerate(experiments):
        for j, each_matrix_size in enumerate(matrix_size):
            time_start = time.time()
            each_experiment(each_matrix_size, repetitions=100)
            d_t = time.time() - time_start
            print(f"{each_matrix_size:d} {each_experiment.__name__:<30s}: {d_t:f}")
            results[i].append(d_t)

    for each_experiment, each_result in zip(experiments, results):
        pyplot.plot(matrix_size, each_result, label=each_experiment.__name__)

    pyplot.legend()
    pyplot.show()


if __name__ == "__main__":
    main()

这是结果（针对矩阵边长的运行时间）：

所以现在，我会坚持使用torch.solve. 然而，最初的问题仍然存在：

如何利用 GPU 更快地求解线性方程？

Answer 1

您的分析在多个方面都是正确的，但有一些细微差别可能有助于澄清您的结果并提高 GPU 性能：

1. CPU 与 GPU 性能 一般来说，GPU 操作会产生与在 CPU 和 GPU 内存之间传输数据相关的开销成本。因此，对于较大的数据集，GPU 加速的好处通常会变得明显，其中并行化的好处超过了这种开销。这种开销成本可能是 GPU 计算对于小矩阵较慢的原因。要利用 GPU 求解线性方程，您应该关注更大的矩阵。

2. Torch solvevstorch.linalg.solve 该torch.solve函数自 PyTorch 1.7.0 起已弃用。您可能会获得更好的性能和更准确的结果torch.linalg.solve。

3. Numba 的njit性能 Numba 的@njit装饰器通过在导入时使用 LLVM 编译器基础结构生成优化的机器代码来加速 Python 函数。当您使用@njit装饰器时，Numba 在非 Python 模式下编译函数，如果函数无法完全优化，则可能会导致性能下降。第一次运行还将包括“编译”步骤，如果将其包含在计时中，则可能会使其显得慢得多。

4. 高效使用 CUDA 内存torch.set_default_tensor_type("torch.cuda.FloatTensor")函数中的 行将solve_torch_gpu默认张量类型设置为 CUDA 张量。之后创建的每个张量都将是 CUDA 张量。这可能会导致不必要地使用 GPU 内存并减慢计算速度。如果您在需要时直接在 GPU 上创建张量（使用.to(device)您device的 CUDA 设备），它将更加高效，并且可能会缩短您的计算时间。

这是函数的修订版本，它使用torch.linalg.solve并直接在 GPU 上创建张量：

def solve_torch_gpu_optimized(dim: int = 5):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    a = torch.rand(dim, dim, device=device)
    b = torch.rand(dim, 1, device=device)

    for _ in range(1000):
        torch.linalg.solve(a, b)