PyTorch로 Numpy를 대체하여 성능을 끌어올려보자 #pytorch #numpy #gpu #cuda #python #colab

PyTorch의 Tensor와 Numpy의 ndarray는 유사한 형태를 가지고 있고 PyTorch의 경우 GPU를 사용한 연산이 가능하기 때문에 Numpy로 작업시 연산 부분을 PyTorch대체해서 처리 속도를 끌어 올릴 수 있다.

 

다음과 같이 Numpy로 배열을 만들고 10000번의 연산을 했을때 소요 시간은 5.11초이다.

%%time

import numpy as np

x = np.ones(1024000)
for i in range(10000):
  x += 1
print (type(x))
print (x)

<class 'numpy.ndarray'>

[10001. 10001. 10001. ... 10001. 10001. 10001.]

CPU times: user 5.11 s, sys: 2.45 ms, total: 5.11 s

Wall time: 5.12 s

 

동일한 로직에 Numpy 배열을 Torch Tensor로 변환하고 Tensor를 CPU에서 GPU로 옮기고 연산이 끝난 후에 Tensor를 다시 GPU에서 CPU로 옮긴 후 Tensor를 Numpy 배열로 변환해주는 코드를 추가 했을때 소요 시간은 0.339이다. 이리저리 변환작업이 추가되지만 Numpy로 연산을 했을때보다 Torch로 GPU 연산을 했을때 16배정도 빠르다는 것을 확인할 수 있다.

%%time

import numpy as np
import torch

x = np.ones(1024000)
x = torch.from_numpy(x)
x = x.to(torch.device('cuda'))
for i in range(10000):
  x += 1
b = x.to(torch.device('cpu'))
c = b.numpy()
print (type(c))
print (c)

<class 'numpy.ndarray'>

[10001. 10001. 10001. ... 10001. 10001. 10001.]

CPU times: user 268 ms, sys: 68 ms, total: 336 ms

Wall time: 339 ms

 

Google CoLab에서 주피터 노트북으로 테스트하였다.