Почему PyTorch медленнее, чем PyOpenCL, который медленнее, чем Numba на GPU? - PullRequest
Новая
       96

Почему PyTorch медленнее, чем PyOpenCL, который медленнее, чем Numba на GPU?

2 голосов
/ 25 апреля 2020

Я работал над программой FDTD, в которой использовался дискретный лапласиан, который я могу реализовать как операцию свертки. Из того, что я прочитал , основным компонентом PyTorch является тензорная библиотека, оптимизированная для выполнения операций, обычно используемых для машинного обучения (таких как свертки). Мне было интересно сравнить его с другими фреймворками, которые я использовал, поэтому я написал тестовую программу для многократного применения дискретного лапласиана к массиву 1d и сравнения времени выполнения: 

import torch as tr
import time
from numba import jit, cuda
import numpy as np
import pyopencl as cl
from pyopencl import array


#parameters
number_of_timesteps = 1000
number_of_elements = 10000000


#set up the inital conditions
torch_data = tr.rand((1,1,number_of_elements),dtype=tr.double) #torch convolution needs shape (minibatch,in_channels,iW)
numba_data = np.array([0] + list(torch_data[0][0].numpy()) + [0]) #add padding [0] for convolution. handled automatically in torch.
opencl_data = np.array([0] + list(torch_data[0][0].numpy()) + [0])


#test Torch
device = "cuda"
torch_data_a = torch_data.to(device)
torch_data_b = torch_data.to(device)
kernel = tr.tensor([[[1,-2,1]]],dtype=tr.double,device=device)
with tr.no_grad():
    start_time = time.time()
    for t in range(round(number_of_timesteps/2)): # /2 because each loop is two convolutions
        torch_data_b = torch_data_a + 0.1* tr.nn.functional.conv1d(torch_data_a,kernel,padding=1)
        torch_data_a = torch_data_b + 0.1* tr.nn.functional.conv1d(torch_data_b,kernel,padding=1)
    print("Torch GPU time:",time.time()-start_time)
    torch_data_numpy = np.array([0] + list(torch_data_a[0][0].cpu().numpy()) + [0])

#Numba GPU kernel
@cuda.jit
def numba_conv_cuda(x,x_new):
    gid = cuda.grid(1)
    if 0 < gid < x.size - 1 :  # Check array boundaries
        x_new[gid] = x[gid] + 0.1*(x[gid+1]+x[gid-1]-2*x[gid])

threadsperblock = 100
blockspergrid = (numba_data.size + (threadsperblock - 1)) // threadsperblock
x_a = cuda.to_device(numba_data)
x_b = cuda.to_device(numba_data)
start_time = time.time()
#actually run the kernel
for t in range(round(number_of_timesteps/2)): #again /2 because each loop is two convolutions
    numba_conv_cuda[blockspergrid, threadsperblock](x_a,x_b)
    numba_conv_cuda[blockspergrid, threadsperblock](x_b,x_a)
print("Numba GPU time:",time.time()-start_time)
numba_data = x_a.copy_to_host()


#test OpenCL
context = cl.create_some_context(interactive=False,answers=[0])
queue = cl.CommandQueue(context)
mem_flags = cl.mem_flags
program = cl.Program(context, """
    #pragma OPENCL EXTENSION cl_khr_fp64 : enable //enable double precision calculations
    __kernel void update_psi(__global const double *x, __global double *x_new)
    {
        int gid = get_global_id(0);
        if(0 < gid && gid < x.size - 1){
            x_new[gid] = x[gid] + 0.1*(x[gid+1]+x[gid-1]-2*x[gid]);
        }
    }
    """.replace("x.size",str(opencl_data.size))).build()
x_a_buf = cl.Buffer(context, mem_flags.READ_WRITE | mem_flags.COPY_HOST_PTR, hostbuf=opencl_data)
x_b_buf = cl.Buffer(context, mem_flags.READ_WRITE | mem_flags.COPY_HOST_PTR, hostbuf=opencl_data)


#actually run the OpenCL
start_time = time.time()
for t in range(round(number_of_timesteps/2)): #again /2 because each loop is two convolutions
    event = program.update_psi(queue, [threadsperblock*blockspergrid], [threadsperblock], x_a_buf, x_b_buf)
    event.wait()
    event = program.update_psi(queue, [threadsperblock*blockspergrid], [threadsperblock], x_b_buf, x_a_buf)
    event.wait()
print("OpenCL GPU time:",time.time()-start_time)
event = cl.enqueue_copy(queue, opencl_data, x_a_buf)
event.wait()


print("Results are same?",np.allclose(torch_data_numpy,numba_data) and np.allclose(numba_data,opencl_data))

И вот результаты тестирования на графическом процессоре Nvidia: 

Torch GPU time: 13.544365406036377
Numba GPU time: 0.2404193878173828
OpenCL GPU time: 0.9025869369506836
Results are same? True

Я удивлен, что результаты показывают, что библиотека, предназначенная для применения таких операций, как свертки, должна быть намного медленнее, чем Numba или PyOpenCL (что даже не оптимизировано, потому что я этого не делал использовать любую локальную память на GPU). Это действительно так, или я что-то не так сделал?

Кроме того, почему ядро, записанное в c, более чем в 3 раза медленнее, чем ядро, написанное в Python?

...