Я пытаюсь решить множество линейных уравнений как можно быстрее. Чтобы найти самый быстрый способ, я протестировал NumPy и PyTorch , каждый на ЦП и на моем графическом процессоре GeForce 1080 (используя 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
Самый медленный - это PyTorch с ускорением CUDA. Я подтвердил, что PyTorch использует мой графический процессор с
import torch
torch.cuda.is_available()
torch.cuda.get_device_name(0)
, возвращающим
True
'GeForce GTX 1080'
Я могу отстать от этого, на ЦП PyTorch медленнее, чем NumPy. Я не могу понять, почему PyTorch на GPU работает намного медленнее. Не так важно, но на самом деле еще больше сбивает с толку то, что декоратор njit в Numba снижает производительность на порядок, до тех пор, пока вы больше не будете использовать синтаксис @ decorator.
Это это моя установка? Иногда я получаю странное сообщение о том, что файл страницы / подкачки windows недостаточно велик. Если я выбрал совершенно непонятный путь к решению линейных уравнений на GPU, я был бы счастлив, если бы меня направили в другом направлении.
Edit
Итак, я сосредоточился на 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()
И вот результаты (время выполнения в зависимости от длины края матрицы):
Edit 2
Видя, что Numba не имеет большого значения в моем случае, я вернулся к тестированию PyTorch. И действительно, он выглядит примерно в 4 раза быстрее, чем Numpy, даже без использования устройства 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. Однако остается исходный вопрос:
Как я могу использовать свой графический процессор, чтобы решать линейные уравнения еще быстрее?