AI
TensorFlow Custom Op으로 데이터 변환 최적화하기
두줄요약
TensorFlow 기반 N-gram 마스킹 위치 계산 병목을 C++ Custom Op으로 교체했습니다. 문장당 처리 시간을 4.697ms에서 17.147μs로 줄이며 약 250배 성능을 개선했습니다.
문제 상황
- 대규모 코퍼스의 TFRecord 변환 과정에서 N-gram 마스킹 위치 계산이 전체 시간의 절반 이상 차지
- TPU 기반 언어 모델 학습용 BERT 입력, 마스킹 위치, NSP·SOP 레이블 생성 요구
구조와 흐름
- TensorFlow
while_loop와random.shuffle기반 마스킹 후보 생성·셔플 방식 - C++
REGISTER_OP,OpKernel,REGISTER_KERNEL_BUILDER로 CPU Custom Op 구현 - g++로 Shared Object 컴파일 후
tf.load_op_library로 Python 코드에서 로딩
성능/운영 포인트
- 문장당 처리 시간 4.697ms에서 17.147μs로 감소
- 동일 기능 기준 약 250배 빠른 마스킹 위치 계산
주의할 점
- Bazel 미사용 환경의 유지보수 난이도
- Op 이해, 디버깅, 구현 복잡도 증가
- 병목이 명확하고 기존 Op 표현이 비효율적인 구간에 한정한 적용
