TF Keras 학습 속도 줄이기 (30%)

라이브러리/Tensorflow keras

rongxian 2022. 4. 26. 10:15

만약 최신 GPU로 TF Keras를 이용하고 있다면, 몇 줄만 추가해서 30% 정도 학습속도를 줄일 수 있다.

from tensorflow.keras import mixed_precision
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)

*주의: TF에선 마지막 레이어가 softmax (sigmoid인 경우도)라면 stability를 위해 float32로 바꾸는 걸 추천함!

output_layer = tf.keras.layers.Activation('sigmoid', dtype='float32')(output_layer)