pytorch에서 tensor 타입 변경에 따른 training time과 memory의 trade-off

sw.kim 2025. 4. 30. 22:03

제목에 영어가 참 많다.. 이러기 싫은데, 억지스럽지 않게 한글로 표현하기가 너무 어렵다.

우선 해당 상황은 pytorch를 이용하여 복소수 신경망(CVNN, CVCNN)을 구현하는 중 발생했다.
complex를 사용하는 경우가 아니라면 아래와 같은 형 변환을 할 일도, 이유가 많지 않기 때문에 이 trade off는 좀 특수한 케이스일 것이라고 생각된다. 다만 이 문제로 많은 고생을 하였기에 혹시나 싶은 마음에 아래 글을 남겨둔다.

복소수 신경망 모델을 구현하고 학습하는 과정에서 GPU 메모리 피크 사용량이 과도하게 증가하는 문제가 발생하였다. 이에 대한 원인을 분석하고 코드 차원에서 메모리 효율을 개선하기 위한 작업을 수행하였다. 결과적으로 메모리 사용량은 줄어들었지만, 학습 시간이 증가하는 현상이 발생하였고, 이에 대한 분석 결과를 정리한다.

문제 분석 (높은 메모리 사용량)

초기 구현된 복소수 Convolution, Linear, BatchNorm 레이어는 복소수 입력 텐서를 실수부(real)와 허수부(imag)로 분리한 후, 각각에 대해 독립적으로 연산을 수행하는 방식이었다. 이때 실수부와 허수부 연산 결과는 각각 float32 텐서로 생성되었으며, 최종 복소수 출력은 torch.complex(out_r_float, out_i_float) 함수를 통해 생성되었다.

# 초기 구현 방식 (예시)
x_r = x_complex.real  # float32 텐서
x_i = x_complex.imag  # float32 텐서

# 실수/허수 부분 계산 수행, 결과는 float32 텐서
out_r_float = calculate_real_part_ops(x_r, x_i, weights_r, weights_i)
out_i_float = calculate_imag_part_ops(x_r, x_i, weights_r, weights_i)

# 최종 복소수 텐서 생성 (문제 발생 지점)
complex_output = torch.complex(out_r_float, out_i_float)

PyTorch Autograd 시스템은 역전파에 필요한 기울기 계산을 위해 순전파 중 생성된 중간 텐서를 메모리에 유지한다. torch.complex 함수는 기존의 out_r_float와 out_i_float 텐서를 입력으로 새로운 complex_output 텐서를 생성하며, Autograd는 이 복소수 텐서에 대한 기울기 계산을 위해 두 float32 텐서의 정보를 계속 유지하게 된다. 그 결과 동일한 정보가 실수부, 허수부, 그리고 복소수 형태로 총 세 번 메모리에 상주하게 되어 피크 메모리 사용량이 급증하게 된다.

해결책 (코드 수정 내용)

메모리 사용량을 줄이기 위해 코드 수정이 이루어졌다. 실수부와 허수부 계산 방식은 기존과 동일하게 유지하되, 최종 복소수 텐서를 생성하는 방식을 변경하였다. torch.complex() 대신, 각각의 float32 텐서를 명시적으로 complex64 타입으로 캐스팅한 후, 복소수 덧셈을 통해 결합하는 방식을 사용하였다.

# 수정된 구현 방식 (예시)
x_r = x_complex.real  # float32 텐서
x_i = x_complex.imag  # float32 텐서
dtype = x_complex.dtype  # 입력 텐서의 complex dtype (예: torch.complex64)

# 실수/허수 부분 계산 수행, 결과는 float32 텐서
out_r_float = calculate_real_part_ops(x_r, x_i, weights_r, weights_i)
out_i_float = calculate_imag_part_ops(x_r, x_i, weights_r, weights_i)

# float 결과를 complex dtype으로 즉시 캐스팅 후 복소수 덧셈으로 결합
complex_output = out_r_float.type(dtype) + 1j * out_i_float.type(dtype)

이 방식은 float32 결과 텐서를 complex64로 변환하는 시점을 명시적으로 지정해 줌으로써, Autograd가 float32 중간 텐서들을 더 빠르게 메모리에서 해제할 수 있도록 유도한다. 그 결과 동시에 메모리에 존재하는 텐서 수가 줄어들어 피크 메모리 사용량을 효과적으로 감소시킬 수 있다.

결과

수정된 코드 적용 후 GPU 피크 메모리 사용량이 눈에 띄게 감소하였다. 구체적으로는 22GB에서 14.5GB로 약 2/3으로 감소하였으며 이로 인해 기존 약 500이 한계였던 batch size를 1000까지 늘릴 수 있었다.

그러나 학습 시간은 기존 98.8초에서 109.7초로 약 10% 정도 증가하였다.

학습 시간 증가 원인 분석

메모리 사용량은 성공적으로 줄였지만, 학습 시간 증가는 아래와 같은 연산 및 처리 오버헤드로 인해 발생한 것으로 보인다.

명시적 타입 캐스팅 오버헤드
.type(dtype) 호출은 기존 텐서를 새로운 타입으로 복사 또는 재구성하는 연산이다. torch.complex()가 내부적으로 최적화된 방식으로 복소수를 생성하는 반면, 명시적 타입 캐스팅은 각 연산마다 추가적인 연산 비용을 유발하고, 이는 누적되어 학습 시간 증가로 이어진다.
복소수 덧셈 연산 오버헤드
tensor.type(...) + 1j * tensor.type(...) 방식은 실제 복소수 연산을 수행하게 된다. 이는 단순히 메모리 상에서 데이터를 결합하는 것보다 복잡한 연산이며, 그만큼 시간이 더 소요될 수 있다.
최적화 수준의 차이
PyTorch 및 하위 CUDA 라이브러리는 float → float → torch.complex와 같은 전형적인 연산 패턴에 대해 이미 고도로 최적화되어 있을 수 있다. 반면, 새롭게 적용된 연산 방식은 비교적 비표준적인 접근이기 때문에 하드웨어 및 라이브러리 차원에서 최적화가 덜 이루어졌을 가능성이 있다.

이러한 복합적인 이유로, 연산 흐름이 변경되면서 새로운 패턴에서 연산 오버헤드가 발생하고 전체 학습 시간 증가로 이어진 것으로 분석된다.

결론

이번 코드 수정은 복소수 신경망 모델에서 GPU 메모리 사용량을 줄이기 위한 목적을 효과적으로 달성하였다.

다만, 메모리 사용량 감소와 맞바꿔 학습 시간이 약 10% 증가하는 트레이드오프가 발생하였다.