Train Faster: 텐서플로우 성능 최적화 기법

GTC 2021에서 발표되었던 Train Faster: A Guide to Tensorflow Performance Optimization를 정리하였다.

1. Automatic Mixed Precision

• 학습에 일반적으로 FP32를 사용하나 Conv나 MatMul 연산 등은 FP16을 사용
• 어떠한 operation들은 FP32를 사용함
• 같은 하이퍼-파라미터를 사용하여 FP32 학습과 같은 정확도를 달성
• 메모리 foot-print를 줄이고 memory-intensive operations를 가속함

Code Snippets

policy = tf.keras.mixed_precision.Policy(‘mixed_float16’)

policy = tf.keras.mixed_precision.Policy(‘mixed_bfloat16’) # TPU

tf.keras.mixed_precision.set_policy(policy)

2. NCHW vs NHWC

• Conv 연산 시 NHWC가 NCHW에 비해 더 빠른 데이터 포맷임
• Tensor Core Utilization를 개선
• 불필요한 Data Layout Transposes를 제거함

Code Snippets

tf.keras.layers.Conv2D(…., data_format=”channels_last”)

3. Auto-clustering

• 왜 한 번에 하나의 노드만을 실행해야 할까? XLA를 이용하여 커널을 퓨전하자!
• Note: XLA는 모든 shape를 위한 새로운 커널을 컴파일해야 하기 때문에 동적 shapes를 모델(e.g., transformer)에 잘 동작하지 않는다.

Code Snippets

$ TF_XLA_FLAGS=” — tf_xla_auto_jit=<level>”

level: -1(disable), 1,2 (Higher level = XLA is more aggressive)

4. Bias

• 필요하지 않으면 Bias를 제거함
• Batch norm는 평균에 의해 값을 이동시키고, 이 이동은 상수 bias항의 필요성을 제거함

5. Batch Size

• Batch Size를 높일수록 학습 속도가 늘어남
• Batch size를 증가시키는 것은 모델 정확도를 변화하므로 타겟 정확도를 만족하기 위해 learning rate와 같은 하이퍼-파라미터를 튜닝이 필요함

메모리 사용률을 확인하기 위해 nvidia-smi을 사용

nvidia-smi — query-gpu=utilization.memory,memory.total,memory.free,

memory.used — format=csv

6. Data Pipeline Optimization

• tf.data API를 사용해라!!
• tf.data.Dataset.interleave : 병렬로 TFRecords를 읽기
• num_parallel_calls=tf.data.experimental.AUTOTUNE: 병렬로 operations를 실행함
• dataset.shard(): 미리 정의된 정크 크기로 매우 큰 데이터셋을 슬라이스하는데 사용
• dataset.prefetch(batch_size): 데이터셋을 반환하기 직전에 미리 prefetch함
• dataset.cache(): 전체 데이터셋을 메모리에 들어갈 수 있을때만
• tf.data.Options()