[PyTorch] 파이토치 코드 실행 속도 높이기 / .cuda()와 device=torch.device('cuda')의 차이 + 애플 m1, m2 칩

PyTorch에서 GPU로 작업할 때 GPU 디바이스에 데이터를 넘겨주는 여러 가지 방법이 있다.

# Approach 1
torch.rand(10, 10).cuda()

# Approach 2
torch.rand(10, 10, device=torch.device('cuda'))

 

두 가지 방법의 차이와 속도를 비교해 보자.

---

# Approach 1
torch.rand(10, 10).cuda()

이 방법에서 torch.rand 함수는 CPU에서 임의의 텐서를 생성한다.

그런 다음 .cuda() 메서드가 텐서에서 호출되어 텐서를 CPU에서 GPU로 이동한다.

이 연산은 inplace로 수행되어 원래 텐서를 수정하고 GPU에 동일한 텐서를 반환하는 형태로 작동한다.

 

# Approach 2
torch.rand(10, 10, device=torch.device('cuda'))

반면 이 방법에서 torch.rand 함수는 cuda로 설정된 device 인자(argument)와 함께 호출된다.

이런 경우 PyTorch가 텐서를 CPU를 거치지 않고 직접 GPU 메모리에 텐서를 할당하도록 한다.

 

두 코드는 실행 속도에서 차이가 발생하는데 추가적인 메모리 전송으로 인해 차이가 발생한다.

.cuda()을 사용하는 첫 번째 방식에서는 텐서가 CPU에서 생성되고 GPU 메모리로 이동된다.

이런 두 디바이스 간의 데이터 전달은 오버헤드(overhead: 주요 작업이나 계산과 직접 관련되지 않은 특정 작업 때문에 생기는 추가적인 비용이나 시간)를 발생시키며, 특히 큰 텐서의 경우 전체 실행 시간에 영향을 미칠 수 있다.

반면, 두 번째 접근 방식은 텐서를 직접 GPU 메모리에 할당하여 추가적인 전송하는 단계를 피하고 더 빠른 실행 속도를 가져갈 수 있다.

---

직접 예제를 통해 확인해보자.

import torch
import time

num_runs = 10

# Approach 1: Generate on CPU and move to GPU
execution_times1 = []
for _ in range(num_runs):
    start_time = time.time()
    tensor1 = torch.rand(10000, 10000).cuda()
    end_time = time.time()
    execution_time1 = end_time - start_time
    execution_times1.append(execution_time1)

average_time1 = sum(execution_times1) / num_runs

# Approach 2: Allocate directly on GPU
execution_times2 = []
for _ in range(num_runs):
    start_time = time.time()
    tensor2 = torch.rand(10000, 10000, device=torch.device('cuda'))
    end_time = time.time()
    execution_time2 = end_time - start_time
    execution_times2.append(execution_time2)

average_time2 = sum(execution_times2) / num_runs

print("Approach 1 average execution time:", average_time1)
print("Approach 2 average execution time:", average_time2)

좀더 드라마틱한 비교를 위해 텐서를 1억개의 랜덤한 값을 가지는 텐서로 설정하였다.

각 방법에 대한 실행 시간을 측정하고 인쇄하는 방식인데 좀더 객관적인 비교를 위해 10번 실행하고 각 시간을 평균내는 방법을 선택했다.

코드를 직접 실행하여 결과를 비교해보자.

두 번째 방법이 빠른 것을 확인할 수 있을 것이다.

 

참고로 내 결과는 이렇다.

Approach 1 average execution time: 1.0137641429901123

Approach 2 average execution time: 0.00017817020416259765

---

Apple Silicon 칩에서 비교해보기

import torch
import time

num_runs = 10

# Approach 1: Generate on CPU and move to MPS
execution_times1 = []
for _ in range(num_runs):
    start_time = time.time()
    tensor1 = torch.rand(10000, 10000).to('mps')
    end_time = time.time()
    execution_time1 = end_time - start_time
    execution_times1.append(execution_time1)

average_time1 = sum(execution_times1) / num_runs

# Approach 2: Allocate directly on MPS
execution_times2 = []
for _ in range(num_runs):
    start_time = time.time()
    tensor2 = torch.rand(10000, 10000, device=torch.device('mps'))
    end_time = time.time()
    execution_time2 = end_time - start_time
    execution_times2.append(execution_time2)

average_time2 = sum(execution_times2) / num_runs

print("Approach 1 average execution time:", average_time1)
print("Approach 2 average execution time:", average_time2)

이 코드로 m1, 2칩에서 비교해보자.

첫 번째 방법에서 .cuda()을 .to('mps'), 두 번째 방법에서 cuda자리를 mps로 바꾸면 된다.

 

m1 pro에서 결과는 아래와 같았다.

Approach 1 average execution time: 0.35477118492126464

Approach 2 average execution time: 0.00010352134704589843

주목할만한 점은 cpu에서 mps 디바이스로 넘기는 첫 번째 과정이 좀더 빨랐다는 것이다.