[PyTorch] Tensor 연산할 때 t1**2가 빠를까? t1.pow(2)가 빠를까?

PyTorch로 코딩을 하다 보면 같은 연산임에도 불구하고 측정할 때마다 속도가 다르게 나오거나, 특정 방식이 미세하게 더 빠를 때가 있다.
이번글에서는**2 연산자와 .pow(2) 메서드 사이의 속도 차이가 발생하는 근본적인 이유를 정리하려고 한다.

---

문제상황

pytorch에서 텐서에 대해 제곱을 계산할때,
**2과 .pow(2)중 어떤것이 더 빠른지 궁금해서 간단한 실험을 준비했다.

• 환경 : colab T4

t1 = torch.arange(start = 0, end = 1000, step = 0.001)
t1 = t1.to(device="cuda")

start_time = time.time()
print(t1**2)
end_time = time.time()
print(f"t1**2 : {(end_time - start_time)*1000}ms")

start_time = time.time()
print(t1.pow(2))
end_time = time.time()
print(f"t1.pow(2): {(end_time - start_time)*1000}ms")

start_time = time.time()
print(t1.pow(2))
end_time = time.time()
print(f"t1.pow(2): {(end_time - start_time)*1000}ms")

start_time = time.time()
print(t1**2)
end_time = time.time()
print(f"t1**2 : {(end_time - start_time)*1000}ms")

tensor([0.0000e+00, 1.9881e-02, 7.9524e-02,  ..., 1.9881e+10, 1.9881e+10,
        1.9881e+10], device='cuda:0')
t1**2 : 1.4524459838867188ms
tensor([0.0000e+00, 1.9881e-02, 7.9524e-02,  ..., 1.9881e+10, 1.9881e+10,
        1.9881e+10], device='cuda:0')
t1.pow(2): 0.904083251953125ms
tensor([0.0000e+00, 1.9881e-02, 7.9524e-02,  ..., 1.9881e+10, 1.9881e+10,
        1.9881e+10], device='cuda:0')
t1.pow(2): 0.8094310760498047ms
tensor([0.0000e+00, 1.9881e-02, 7.9524e-02,  ..., 1.9881e+10, 1.9881e+10,
        1.9881e+10], device='cuda:0')
t1**2 : 0.7979869842529297ms

왜 수학적으로는 동일한 연산임에도 불구하고 시간차이가 나타나는 걸까?

---

1. CUDA Warm-up & Synchronization

GPU 연산을 처음 시작할 때 가장 많이 겪는 현상이다.

• Warm-up: GPU는 첫 연산 시 컨텍스트 초기화, 커널 로드 등의 준비 시간이 필요하다. 따라서 코드 상에서 먼저 실행되는 연산이 이 "준비 시간"을 독박 쓰게 되어 훨씬 느려 보일 수 있다.
• 비동기 실행: GPU 연산은 CPU에서 명령만 내리고 따로 돌아간다. print()를 통해 결과를 확인하려 하면, CPU는 GPU가 연산을 끝내고 데이터를 보내줄 때까지 기다려야 하므로, 이 대기 시간이 측정값에 포함된다.

이를 반영하여 아래처럼 코드를 짰다.

# Warm-up (미리 한번 실행해서 GPU를 깨워둠)
_ = t1**2 
torch.cuda.synchronize() # GPU 작업이 끝날 때까지 대기

# t1**2 측정
start = time.time()
res1 = t1**2
torch.cuda.synchronize() # 연산이 완전히 끝날 때까지 기다린 후 시간 측정
print(f"t1**2: {(time.time() - start)*1000:.4f}ms")

# t1.pow(2) 측정
start = time.time()
res2 = t1.pow(2)
torch.cuda.synchronize()
print(f"t1.pow(2): {(time.time() - start)*1000:.4f}ms")

# t1.pow(2) 측정
start = time.time()
res2 = t1.pow(2)
torch.cuda.synchronize()
print(f"t1.pow(2): {(time.time() - start)*1000:.4f}ms")

# t1**2 측정
start = time.time()
res1 = t1**2
torch.cuda.synchronize()
print(f"t1**2: {(time.time() - start)*1000:.4f}ms")

t1**2: 0.0567ms
t1.pow(2): 0.0527ms
t1.pow(2): 0.0515ms
t1**2: 0.0520ms

이번에는 아까처럼 눈에 띄는 정도의 시간차는 보이지 않아졌다. 그럼에도 불구하고 미세한 시간차이가 나는 이유가 무엇일까?

2. 파이썬 인터프리터의 오버헤드

Warm-up을 제외하더라도 마이크로초($\mu s$) 단위의 미세한 차이가 지속된다면, 이는 파이썬의 처리 경로 차이 때문이다.

• t1**2 (연산자 오버로딩):
  1. 파이썬 인터프리터가 ** 연산자 인식.
  2. 객체의 __pow__ 매직 메서드를 찾는 과정을 거침.
  3. 최종적으로 PyTorch의 C++ 백엔드 호출.
• t1.pow(2) (메서드 직접 호출):
  1. 객체에 정의된 pow 메서드를 직접 호출.
  2. 연산자 해석 단계를 건너뛰므로 함수 호출 오버헤드(Function Call Overhead)가 더 적다.

3. 커널 런칭(Kernel Launch)의 효율성

GPU 연산 속도는 '계산 시간'보다 '명령 전달 시간'에 좌우되기도 한다.

• t1.pow(2)는 PyTorch의 C++ API로 진입하는 경로가 더 직접적이다.
• CPU가 GPU에게 "이 연산을 실행해!"라고 명령을 전달하는 과정이 아주 미세하게 더 효율적일 수 있다.

4. 시스템 노이즈와 캐싱

측정 시마다 값이 튀는 것은 컴퓨터의 하드웨어 상황 때문이다. CPU 스케줄링이나 명령어 캐시 상태에 따라 0.01ms 정도의 오차는 언제든 발생할 수 있으며, 두 번째 실행부터는 이미 캐시에 로드된 정보를 사용하기 때문에 더 빨라지기도 한다.

---

결론: 무엇을 써야 할까?

항목	t1**2	t1.pow(2)
내부 구현	연산자 오버로딩	C++ API 직접 호출
속도	미세하게 느릴 수 있음	미세하게 빠를 수 있음
추천 용도	직관적인 수학 수식 표현	일관된 함수형 프로그래밍 스타일

결론적으로, 성능 차이는 실무에서 무시해도 될 만큼 미미해보인다.
이런경우, 성능 최적화보다는 팀의 컨벤션이나 코드의 가독성을 우선시하여 선택하는 것이 좋다.
만약 정확한 성능 비교를 원한다면 반드시 torch.cuda.synchronize()를 사용하여 GPU의 작업이 끝날 때까지 기다린 후 시간을 측정해야 한다.