我正试图尽快解决大量的线性方程组。为了找出最快的方法，我对NumPy和PyTorch进行了基准测试，它们分别在CPU和GeForce 1080 GPU上（NumPy使用Numba）。结果真的让我很困惑。
这是我在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：高达500 x500的矩阵大小，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的基准测试。事实上，它似乎比Numpy快4倍，甚至没有使用CUDA设备。
下面是我使用的代码：

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更快地求解线性方程组？