하나의 조직에서 TensorFlow와 PyTorch 동시 활용하기

리서치는 PyTorch로 유연하게 수행하고 프로덕션 배포와 대형 학습은 TensorFlow로 가져가는 이중 프레임워크 운영기다. 같은 구조의 모델을 두 프레임워크로 구현하고 checkpoint와 weight를 양방향 변환하는 내부 라이브러리를 만들었다. Dense weight 전치, GRU gate 순서, 부동소수점 오차처럼 결과 동등성을 깨뜨리는 세부 차이와 테스트 방법까지 다룬다.

핵심 포인트

PyTorch는 연구 코드 변경에 유리했지만 당시 Python·HTTP overhead 없는 표준 배포 경로가 부족해 TensorFlow Serving을 선택했다.
TensorFlow SavedModel은 모델 파일과 Docker image만으로 C++ 기반 gRPC·HTTP 서버를 띄울 수 있었다.
내부 모델을 TensorFlow와 PyTorch로 각각 구현하고 checkpoint, state dict, live model 사이 weight 변환기를 작성했다.
변환 후 동일한 numpy 입력을 두 모델에 넣고 `np.allclose`로 출력 동등성을 반복 검증했다.
TensorFlow Dense와 PyTorch Linear는 weight shape가 반대이며 GRU는 reset·update gate의 배열 순서도 다르다.
코드 이중화 비용은 크지만 연구 편의, TPU 학습, tf.data pipeline, 안정적인 serving을 동시에 확보했다.

상세 정리

선택 배경: 팀은 TensorFlow에서 PyTorch로 통일했다가 대형 모델 배포가 어려워지자 연산과 weight가 같으면 결과도 같다는 전제 아래 TensorFlow를 다시 도입했다.
Serving 요구: 작은 모델은 Python과 HTTP 처리 시간이 추론보다 길 수 있어 해당 overhead를 제거하고 싶었다. 당시 선택지는 TensorFlow Serving, ONNX, libtorch였고 TorchServe는 아직 없었다.
배포 판단: ONNX는 뚜렷한 운영 도구가 부족했고 libtorch는 별도 C++ 통합이 필요했다. TensorFlow Serving은 SavedModel 하나로 고성능 C++ gRPC·HTTP server를 실행할 수 있었다.
TensorFlow 유지 이점: custom op를 피하면 distribution strategy 변경으로 TPU를 쓸 수 있다. `tf.data.Dataset`은 CPU·GPU·TPU에서 같은 pipeline을 실행하며 대형 학습 데이터 코드를 단순화했다.
이중 코드 구조: 리서치 project의 최종 PyTorch 모델을 읽기 쉬운 engineering library로 다시 만들고 동일 구조의 TensorFlow 모델을 별도 codebase에 구현했다.
변환 방향: PyTorch 연구 weight를 TensorFlow로 옮겨 배포하고, TensorFlow로 학습한 대형 모델을 PyTorch로 넘겨 추가 연구하는 두 방향을 모두 지원했다.
Checkpoint→PyTorch: `tf.train.list_variables`로 변수 목록을 보고 `load_variable`로 numpy array를 얻은 뒤 `torch.from_numpy`로 `Parameter.data`에 넣었다.
State dict→TensorFlow: GPU에서 저장된 PyTorch 파일은 `map_location=cpu`로 읽고 tensor를 `detach().numpy()`로 바꿨다. Keras layer의 `set_weights`가 요구하는 kernel·bias 순서를 맞췄다.
Live model 변환: TensorFlow의 `get_weights()` 결과를 PyTorch parameter에 넣거나, PyTorch `state_dict()`의 weight와 bias를 TensorFlow `set_weights()`에 전달했다.
테스트: 두 layer를 같은 hyperparameter로 build하고 한쪽 weight를 변환한 뒤 100개의 random float32 입력 결과를 비교했다. LayerNorm은 구현 오차를 고려해 `rtol=1e-5`, `atol=1e-6`를 사용했다.
Matrix shape: PyTorch Linear는 `(output, input)`, TensorFlow Dense는 `(input, output)` 형태이므로 transpose하지 않으면 같은 연산이 되지 않는다.
수치 오차: 같은 matrix multiplication도 프레임워크별 절대 오차가 약 `2.86e-6` 발생했다. 깊은 network에서 누적될 수 있어 필요한 정밀도에 맞춰 tolerance를 결정해야 한다.
GRU 차이: TensorFlow는 update, reset, candidate 순서이고 PyTorch는 reset, update, candidate 순서다. input·hidden kernel을 transpose하고 gate 구간을 바꾸며 두 bias를 합치거나 나눠야 한다.
결과: 연구자는 모든 내부 모델의 PyTorch 버전을 자유롭게 쓰고, 확정 weight는 빠르게 TensorFlow로 변환해 serving할 수 있게 됐다. 비용이 큰 대신 연구와 운영의 요구를 한 조직에서 분리해 만족시켰다.