GPU 연산 잘 사용하기

이전에 어디선가 포스팅한 글을 읽었는데 (아마 velog일 것이다)

내가 생각하는 최적화와 딥러닝에서 돌아가는 최적화의 갭이 매우 크다라는 것을 느꼈다.

 

딥러닝하는 사람들이라면 모두 "GPU는 행렬 연산에 최적화되어 있다" 라는 말을 많이 봤을 것이다.

딥러닝에서 사용하는 연산은 행렬에 최적화되어 있지만

 

프로그래밍으로 접하는 사람들은 단순히 반복문을 적게 쓰고 Big-O 을 줄이는 방식으로 많이들 생각하겠지만

GPU의 최적화는 무조건 행렬이다.

 

import torch
import random
import time


RANGE = 10000
device = 'cuda'

data = torch.Tensor([random.random() for _ in range(RANGE)]).to(device)
mul_torch = torch.Tensor([idx for idx in range(RANGE)]).to(device)

# for loop
start_time = time.time()

for _idx, _mul in enumerate(mul_torch):
    data[_idx] *= _mul
    
print('for time:', time.time() - start_time)


# multiply
start_time = time.time()

mul_torch = torch.diag(mul_torch)
data = torch.matmul(data, mul_torch)

print('matmul time:', time.time() - start_time)

 

위 코드를 돌렸을 때 for 문을 수행한 코드는 약 0.1769 초 전후로 걸린다.

matrix 곱을 수행한 코드는 약 0.0021 초 전후로 80배 넘게 차이가 난다.

 

몇 번 돌려도 시간 상 큰 차이가 없어서 matrix 곱 연산이 압도적으로 시간이 적게 드는 것을 알 수 있다.

 

앞으로 deep learning 코드에서 for 문이 있다면 행렬로 바꿀 수 없는지 고민해볼 만한 요소로 생각된다.