mixed_precision.ipynb

Kernel: Python 3

Copyright 2019 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

혼합 정밀도

TensorFlow.org에서 보기

Google Colab에서 실행

GitHub에서 소스 보기

노트북 다운로드하기

개요

혼합 정밀도는 훈련 중에 모델에서 16비트 및 32비트 부동 소수점 유형을 모두 사용하여 모델을 더 빠르게 실행하고 메모리를 더 적게 사용하는 데 목적을 두고 있습니다. 수치 안정성을 위해 모델의 특정 부분을 32비트 유형으로 유지하면 모델의 단계 시간이 짧아지고 정확도와 같은 평가 지표에서도 동일하게 잘 훈련할 수 있습니다. 이 가이드에서는 Keras 혼합 정밀도 API를 사용하여 모델 속도를 높이는 방법을 설명합니다. 이 API를 사용하면 최신 GPU에서는 3배 이상, TPU에서는 60% 이상, 최신 인텔 CPU에서는 2배 이상의 성능을 높일 수 있습니다.

오늘날 대부분의 모델은 32-bit 메모리를 사용하는 float32 dtype을 사용합니다. 그러나 정밀도가 낮은 dtype인 float16과 bfloat16도 있으며 각각은 16-bit의 메모리를 사용합니다. 최신 가속기는 16-bit 계산을 실행할 수 있는 특수한 하드웨어가 있어 16-bit dtype을 메모리에서 더 빨리 읽을 수 있으므로 16-bit dtype에서 연산을 더 빠르게 실행할 수 있습니다.

NVIDIA GPU는 float32보다 float16에서 연산을 더 빠르게 수행할 수 있으며, TPU 및 지원되는 Intel CPU는 float32보다 bfloat16에서 연산을 더 빠르게 수행할 수 있습니다. 따라서 이러한 정밀도가 낮은 dtype은 해당 장치에서 가능할 때마다 사용해야 합니다. 그러나 변수와 몇 가지 계산은 여전히 수치상의 이유로 float32를 사용해야 모델이 동일한 품질로 훈련할 수 있습니다. Keras 혼합 정밀도 API를 사용하면 float16 또는 bfloat16과 float32를 함께 사용할 수 있으므로 float16/bfloat16의 성능 이점과 float32의 수치 안정성 이점을 얻을 수 있습니다.

참고: 이 가이드에서 "수치 안정성"이라는 용어는 고정밀 dtype 대신 저정밀 dtype을 사용할 때 모델의 품질에 어떤 영향을 미치는지를 나타냅니다. float16 또는 bfloat16에서 "수치적으로 불안정"하다는 것은 이러한 dtype 중 하나에서 실행할 때 float32에서 연산을 실행할 때와 비교하여 모델의 평가 정확성이나 기타 메트릭이 떨어지는 경우를 말합니다.

설치하기

import tensorflow as tf

from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import mixed_precision

지원하는 하드웨어

혼합 정밀도는 대부분의 하드웨어에서 실행되지만, 최신 NVIDIA GPU, Cloud TPU 및 최신 인텔 CPU에서는 모델 속도만 높일 수 있습니다. NVIDIA GPU는 float16과 float32의 혼합 사용을 지원하는 반면 TPU와 인텔 CPU는 bfloat16과 float32의 혼합을 지원합니다.

NVIDIA GPU 중에서 계산 기능이 7.0 이상인 유닛은 float16 행렬 곱셈 및 컨볼루션을 가속화하기 위한 Tensor Cores라는 특수 하드웨어 장치를 가지고 있기 때문에 혼합 정밀도로부터 최대의 성능 이점을 얻을 수 있습니다. 구형 GPU의 경우 혼합 정밀도 사용에 따른 연산 성능 이점은 없지만 메모리 및 대역폭 절약으로 일부 속도 향상이 가능합니다. NVIDIA의 CUDA GPU 웹 페이지에서 해당 GPU의 계산 기능을 조회할 수 있습니다. 혼합 정밀도의 이점을 가장 크게 누릴 수 있는 GPU의 예로는 RTX GPU, Titan V 및 V100, A100이 있습니다.

인텔 CPU 중 4세대 인텔 제온 프로세서(코드명 Sapphire Rapids)부터는 AMX 명령어를 사용하여 bfloat16 계산을 가속화할 수 있으므로(Tensorflow 2.12 이상 필요) 혼합 정밀도의 성능 이점을 가장 크게 누릴 수 있습니다.

참고: Google Colab에서 이 가이드를 실행하면 일반적으로 GPU 런타임에 P100이 연결됩니다. P100은 6.0의 계산 기능을 갖추고 있으며, 속도가 크게 향상되지는 않을 것으로 예상됩니다. CPU 런타임에서 실행하는 경우, 런타임에 AMX가 없는 CPU가 있을 가능성이 높으므로 속도가 느려질 수 있습니다.

다음을 통해 GPU 타입을 확인할 수 있습니다. 이 명령은 NVIDIA 드라이버가 설치된 경우에만 존재하므로 다음과 같은 경우 오류가 발생합니다.

!nvidia-smi -L

모든 Cloud TPU는 bfloat16을 지원합니다.

속도 향상이 예상되지 않는 구형 인텔 CPU, AMX가 없는 기타 x86 CPU, 구형 GPU에서도 혼합 정밀도 API는 단위 테스트, 디버깅 또는 API를 시험하는 용도로 계속 사용할 수 있습니다. 그러나 AMX 명령어가 없는 CPU의 mixed_bfloat16과 모든 x86 CPU의 mixed_float16은 상당히 느리게 실행됩니다.

dtype 정책 설정

Keras에서 혼합 정밀도를 사용하려면 일반적으로 dtype 정책이라는 tf.keras.mixed_precision.Policy를 생성해야합니다. Dtype 정책은 레이어가 실행될 레이어를 지정합니다. 이 가이드에서는 'mixed_float16' 문자열로 정책을 구성하고 이를 전역 정책으로 설정합니다. 이로 인해 이후에 생성된 레이어가 float16과 float32를 혼합하여 혼합 정밀도를 사용합니다.

policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)

간단히 말해, 실제로 수행하게 되는 형태로서 set_global_policy 문자열을 직접 전달할 수 있습니다.

# Equivalent to the two lines above
mixed_precision.set_global_policy('mixed_float16')

이 정책은 레이어의 두 가지 중요한 측면, 즉 레이어 계산이 수행되는 dtype과 레이어 변수의 dtype을 지정합니다. 위에서 mixed_float16 정책을 만들었습니다(예, 문자열 'mixed_float16'을 생성자에 전달하여 만든 mixed_precision.Policy). 이 정책에서 레이어는 float16 계산 및 float32 변수를 사용합니다. 계산은 성능을 위해 float16에서 수행되지만, 수치 안정성을 위해서는 변수를 float32로 유지해야 합니다. 이러한 정책 속성을 직접 쿼리 할 수 있습니다.

print('Compute dtype: %s' % policy.compute_dtype)
print('Variable dtype: %s' % policy.variable_dtype)

앞에서 언급했듯이 mixed_float16 정책은 7.0 이상의 계산 기능을 갖춘 NVIDIA GPU에서 성능이 가장 크게 향상됩니다. 이 정책은 다른 GPU 및 CPU에서도 실행되지만 성능이 향상되지 않을 수 있습니다. TPU 및 CPU의 경우 mixed_bfloat16 정책을 대신 사용해야 합니다.

모델 구축

다음으로 간단한 모델을 만들어 봅시다. TensorFlow 런타임의 오버헤드가 일반적으로 실행 시간을 지배하여 GPU의 성능 향상을 무시할 수 있기 때문에 매우 작은 토이 모델은 일반적으로 혼합 정밀도의 이점을 얻지 못합니다. 따라서 GPU를 사용하는 경우 각각 4096개의 유닛으로 두 개의 큰 Dense 레이어를 만들어 봅시다.

inputs = keras.Input(shape=(784,), name='digits')
if tf.config.list_physical_devices('GPU'):
  print('The model will run with 4096 units on a GPU')
  num_units = 4096
else:
  # Use fewer units on CPUs so the model finishes in a reasonable amount of time
  print('The model will run with 64 units on a CPU')
  num_units = 64
dense1 = layers.Dense(num_units, activation='relu', name='dense_1')
x = dense1(inputs)
dense2 = layers.Dense(num_units, activation='relu', name='dense_2')
x = dense2(x)

각 계층에는 정책이 있으며 기본적으로 전역 정책을 사용합니다. 따라서 이전에 전역 정책을 mixed_float16으로 설정했기 때문에 각 Dense 레이어는 mixed_float16 정책을 갖습니다. 이로 인해 밀도가 높은 레이어는 float16 계산을 수행하고 float32 변수를 갖게됩니다. float16 계산을 수행하기 위해 입력을 float16으로 캐스팅하여 결과적으로 출력이 float16이됩니다. 변수는 float32이며 dtype 불일치로 인한 오류를 피하기 위해 레이어를 호출하면 float16으로 캐스팅됩니다.

print(dense1.dtype_policy)
print('x.dtype: %s' % x.dtype.name)
# 'kernel' is dense1's variable
print('dense1.kernel.dtype: %s' % dense1.kernel.dtype.name)

다음으로 출력 예측을 작성합니다. 일반적으로 다음과 같이 출력 예측을 작성할 수 있지만 float16에서는 항상 수치상으로 안정적이지는 않습니다.

# INCORRECT: softmax and model output will be float16, when it should be float32
outputs = layers.Dense(10, activation='softmax', name='predictions')(x)
print('Outputs dtype: %s' % outputs.dtype.name)

모델의 끝에서 softmax 활성화는 float32이어야 합니다. dtype 정책이 mixed_float16이므로, softmax 활성화는 일반적으로 float16 계산 dtype을 가지며 float16 텐서를 출력합니다.

Dense와 softmax 레이어를 분리하고 dtype='float32'를 softmax 레이어에 전달하는 것으로 이 문제를 해결할 수 있습니다.

# CORRECT: softmax and model output are float32
x = layers.Dense(10, name='dense_logits')(x)
outputs = layers.Activation('softmax', dtype='float32', name='predictions')(x)
print('Outputs dtype: %s' % outputs.dtype.name)

dtype='float32'를 softmax 레이어 생성자에 전달하면 레이어의 dtype 정책이 float32 정책으로 재정의되어 계산을 수행하고 변수를 float32로 유지합니다. 마찬가지로 dtype=mixed_precision.Policy('float32')대신 전달할 수도 있습니다. 레이어는 항상 dtype 인수를 정책으로 변환합니다. Activation 레이어에는 변수가 없으므로 정책의 변수 dtype은 무시되지만 float32 정책의 계산 dtype은 softmax 및 모델 출력을 float32로 만듭니다.

모델의 중간에 float16 softmax를 추가하는 것은 괜찮지만 모델 끝의 softmax는 float32로 있어야합니다. 그 이유는 softmax에서 손실까지 흐르는 중간 텐서가 float16 또는 bfloat16이면 숫자 문제가 발생할 수 있기 때문입니다.

float16 계산이 수치상으로 안정적이지 않다고 생각된다면 dtype='float32'을 전달하여 모든 레이어의 dtype을 float32로 재정의할 수 있습니다. 하지만 일반적으로 대부분의 레이어는 mixed_float16 및 mixed_bfloat16으로 충분한 정밀도를 갖기 때문에 모델의 마지막 레이어에서만 필요합니다.

모델이 softmax로 끝나지 않더라도 출력은 여전히 float32이어야 합니다. 이 특정 모델에는 필요하지 않지만, 다음을 사용하여 모델 출력을 float32로 캐스팅 할 수 있습니다.

# The linear activation is an identity function. So this simply casts 'outputs'
# to float32. In this particular case, 'outputs' is already float32 so this is a
# no-op.
outputs = layers.Activation('linear', dtype='float32')(outputs)

그런 다음 모델을 완료 및 컴파일하고 입력 데이터를 생성합니다.

model = keras.Model(inputs=inputs, outputs=outputs)
model.compile(loss='sparse_categorical_crossentropy',
              optimizer=keras.optimizers.RMSprop(),
              metrics=['accuracy'])

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = x_train.reshape(60000, 784).astype('float32') / 255
x_test = x_test.reshape(10000, 784).astype('float32') / 255

이 예제는 입력 데이터를 int8에서 float32로 캐스팅합니다. 255로 나누기가 CPU에 있으며 float16 연산은 float32 연산보다 느리기 때문에 float16으로 캐스팅하지 않습니다. 이 경우 성능 차이는 무시할만하지만, 일반적으로 CPU에서 실행되는 경우 float32에서 입력 처리 계산을 실행해야 합니다. 각 레이어는 부동 소수점 입력을 계산 dtype에 캐스팅하므로 모델의 첫 번째 레이어는 입력을 float16으로 캐스팅합니다.

모델의 초기 가중치가 검색됩니다. 가중치를 로딩하여 처음부터 다시 훈련 할 수 있습니다.

initial_weights = model.get_weights()

Model.fit으로 모델 훈련

다음으로, 모델을 훈련합니다.

history = model.fit(x_train, y_train,
                    batch_size=8192,
                    epochs=5,
                    validation_split=0.2)
test_scores = model.evaluate(x_test, y_test, verbose=2)
print('Test loss:', test_scores[0])
print('Test accuracy:', test_scores[1])

모델이 로그에 스텝당 시간을 출력합니다(예: "25ms/step"). TensorFlow가 모델을 최적화하는 데 시간이 걸리는 첫 epoch 중에는 속도가 느릴 수 있지만 이후에는 스텝당 시간이 안정화됩니다.

Colab에서 이 가이드를 실행하는 경우 혼합 정밀도의 성능을 float32와 비교할 수 있습니다. 이렇게 하려면 "dtype 정책 설정" 섹션에서 정책을 mixed_float16에서 float32로 변경한 다음 이 지점까지 모든 셀을 다시 실행합니다. 계산 기능이 7.X인 GPU에서는 스텝당 시간이 크게 증가하여 혼합 정밀도가 모델 속도를 높이는 것을 확인할 수 있을 것입니다. 가이드를 계속 진행하기 전에 정책을 다시 mixed_float16으로 변경하고 셀을 다시 실행해야 합니다.

최소 8.0(Ampere GPU 이상)의 컴퓨팅 성능을 가진 GPU에서는 float32와 비교하여 혼합 정밀도를 사용할 때 이 가이드의 장난감 모델에서 성능 향상을 볼 수 없습니다. 이것은 tf.linalg.matmul과 같은 특정 float32 연산에서 자동으로 더 낮은 정밀도 산술을 이용하는 TensorFloat-32의 사용 때문입니다. TensorFloat-32는 float32를 사용할 때 혼합 정밀도의 몇 가지 성능 이점을 제공합니다. 그러나 실제 모델에서는 일반적으로, TensorFloat-32가 지원하지 않는 연산 및 메모리 대역폭 절약으로 인해 혼합 정밀도에서 상당한 성능 향상을 볼 수 있습니다.

TPU에서 혼합 정밀도를 실행하는 경우 GPU, 특히 Ampere 이전 GPU에서 혼합 정밀도를 실행하는 것과 비교하여 성능이 크게 향상되지는 않습니다. 이는 TPU가 float32의 기본 dtype 정책의 경우에도 bfloat16에서 특정 연산을 백그라운드에서 수행하기 때문입니다. 이것은 Ampere GPU가 기본적으로 TensorFloat-32를 사용하는 방식과 유사합니다. Ampere GPU와 비교하여 TPU는 일반적으로 실제 모델에서 혼합 정밀도 사용에 따른 성능 향상이 더 적습니다.

많은 실제 모델의 경우, float16 텐서가 메모리의 절반을 차지하므로 혼합 정밀도를 사용하면 메모리 부족없이 배치 크기를 두 배로 늘릴 수 있습니다. 그러나 각 배치가 60,000개의 이미지로 구성된 전체 MNIST 데이터세트로 구성된 모든 dtype에서 모델을 실행할 수 있으므로, 이 토이 모델에는 적용되지 않습니다.

손실 조정

손실 스케일링은 tf.keras.Model.fit이 mixed_float16 정책에서 숫자 언더플로우를 방지하기 위해 자동으로 실행하는 기술입니다. 이 섹션에서는 손실 스케일링이 무엇인지 설명하고 다음 섹션에서는 사용자 정의 훈련 루프에서 이 기술을 사용하는 방법을 설명합니다.

참고: mixed_bfloat16 정책을 사용하는 경우 손실 조정을 수행할 필요가 없습니다.

언더플로우 및 오버플로우

float16 데이터 형식은 float32에 비해 동적 범위가 좁습니다. 이는 $65504$ 이상의 값은 무한대로 오버플로우되고 $6.0 \times 10^{-8}$ 미만의 값은 0으로 언더플로우됨을 의미합니다. float32 및 bfloat16은 동적 범위가 훨씬 높아 오버플로우 및 언더플로우가 문제가 되지 않습니다.

예:

x = tf.constant(256, dtype='float16')
(x ** 2).numpy()  # Overflow

x = tf.constant(1e-5, dtype='float16')
(x ** 2).numpy()  # Underflow

실제로 float16으로 오버플로우가 거의 발생하지 않습니다. 또한 순방향 전달에서 언더플로우가 거의 발생하지 않습니다. 그러나 역방향 전달동안 그래디언트가 0으로 언더플로우 될 수 있습니다. 손실 조정은 이러한 언더플로우를 방지하는 기술입니다.

손실 스케일링 개요

손실 스케일링의 기본 개념은 간단합니다. 손실에 어떤 큰 숫자(예: $1024$)를 곱하면 손실 스케일을 얻을 수 있습니다. 이로 인해 그래디언트가 $1024$만큼 확장되어 언더플로우 가능성을 크게 줄입니다. 최종 그래디언트가 계산되면 $1024$로 나눠서 올바른 값으로 되돌립니다.

이 프로세스의 의사 코드는 다음과 같습니다.

loss_scale = 1024
loss = model(inputs)
loss *= loss_scale
# Assume `grads` are float32. You do not want to divide float16 gradients.
grads = compute_gradient(loss, model.trainable_variables)
grads /= loss_scale

손실 규모를 선택하는 것은 까다로울 수 있습니다. 손실 규모가 너무 낮으면 그래디언트가 여전히 0으로 언더플로우일 수 있습니다. 너무 높으면 반대의 문제가 발생합니다. 그래디언트가 무한대로 오버플로우일 수 있습니다.

이 문제를 해결하기 위해 TensorFlow는 손실 스케일을 동적으로 결정하므로 수동으로 선택할 필요가 없습니다. tf.keras.Model.fit을 사용하면 손실 스케일링이 자동으로 수행되므로 추가 작업을 수행할 필요가 없습니다. 사용자 정의 훈련 루프를 사용하는 경우 손실 스케일링을 위해 특수 옵티마이저 래퍼인 tf.keras.mixed_precision.LossScaleOptimizer를 명시적으로 사용해야 합니다. 이에 대해서는 다음 섹션에서 설명합니다.

사용자 정의 훈련 루프를 사용하여 모델 훈련

지금까지 tf.keras.Model.fit을 사용하여 혼합 정밀도로 Keras 모델을 훈련했습니다. 다음으로, 사용자 정의 훈련 루프와 함께 혼합 정밀도를 사용해 보겠습니다. 사용자 정의 훈련 루프가 무엇인지 아직 모르고 있다면 먼저 사용자 정의 훈련 가이드를 읽어보세요.

혼합 정밀도로 사용자 정의 훈련 루프를 실행하려면 float32에서 실행하는것으로부터 두 가지 변경이 필요합니다.

1. 혼합 정밀도로 모델을 작성합니다 (이미 수행했습니다)

2. mixed_float16 을 사용하는 경우 손실 조정을 명시적으로 사용합니다.

(2)단계에서는 옵티마이저를 래핑하고 손실 조정을 적용하는 tf.keras.mixed_precision.LossScaleOptimizer 클래스를 사용합니다. 옵티마이저와 손실 조정이라는 두 가지 인수가 필요합니다. 동적 손실 규모를 사용하려면 다음과 같이 구성합니다

optimizer = keras.optimizers.RMSprop()
optimizer = mixed_precision.LossScaleOptimizer(optimizer)

원하는 경우 명시적 손실 스케일을 선택하거나 손실 스케일링 동작을 사용자 정의할 수 있지만 알려진 모든 모델에서 잘 작동하는 것으로 확인된 기본 손실 스케일링 동작을 유지하는 것이 좋습니다. 손실 스케일링 동작을 사용자 정의하려면 tf.keras.mixed_precision.LossScaleOptimizer 설명서를 참조하세요.

다음으로, 손실 객체와 tf.data.Dataset을 정의합니다.

loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
train_dataset = (tf.data.Dataset.from_tensor_slices((x_train, y_train))
                 .shuffle(10000).batch(8192))
test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(8192)

다음으로, 훈련 스텝 함수를 정의합니다. 손실 스케일 옵티마이저의 두 가지 새로운 메서드를 사용하여 손실 규모를 조정하고 그래디언트 스케일을 축소합니다.

• get_scaled_loss(loss): 손실에 손실 규모를 곱합니다

• get_unscaled_gradients(gradients): 조정된 그래디언트 목록을 입력으로 취하고 각각을 손실 규모로 나누어 조정을 해제합니다

그래디언트의 언더플로우를 방지하려면 이러한 기능을 사용해야 합니다. LossScaleOptimizer.apply_gradients는 Inf 또는 NaN이 없는 그래디언트를 적용합니다. 이는 또한 손실 규모를 업데이트하는데, 그래디언트에 Inf 또는 NaN이 있으면 손실 규모를 절반으로 줄이고 그렇지 않은 경우 잠재적으로 증가시킵니다.

@tf.function
def train_step(x, y):
  with tf.GradientTape() as tape:
    predictions = model(x)
    loss = loss_object(y, predictions)
    scaled_loss = optimizer.get_scaled_loss(loss)
  scaled_gradients = tape.gradient(scaled_loss, model.trainable_variables)
  gradients = optimizer.get_unscaled_gradients(scaled_gradients)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))
  return loss

LossScaleOptimizer는 훈련 시작 시 처음 몇 단계를 건너뛸 수 있습니다. 최적의 손실 규모를 신속하게 결정할 수 있도록 손실 규모가 크게 시작됩니다. 몇 단계를 거치면 손실 규모가 안정화되고 몇 단계만 건너뜁니다. 이 프로세스는 자동으로 수행되며 훈련 품질에는 영향을 미치지 않습니다.

이제 테스트 스텝을 정의합니다.

@tf.function
def test_step(x):
  return model(x, training=False)

처음부터 다시 학습할 수 있도록 모델의 초기 가중치를 로드합니다.

model.set_weights(initial_weights)

마지막으로, 사용자 정의 훈련 루프를 실행합니다.

for epoch in range(5):
  epoch_loss_avg = tf.keras.metrics.Mean()
  test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
      name='test_accuracy')
  for x, y in train_dataset:
    loss = train_step(x, y)
    epoch_loss_avg(loss)
  for x, y in test_dataset:
    predictions = test_step(x)
    test_accuracy.update_state(y, predictions)
  print('Epoch {}: loss={}, test accuracy={}'.format(epoch, epoch_loss_avg.result(), test_accuracy.result()))

GPU 성능 팁

다음은 GPU에서 혼합 정밀도를 사용할 때의 성능 팁입니다.

배치 크기 늘리기

모델 품질에 영향을 미치지 않으면 혼합 정밀도를 사용할 때 배치 크기를 두 배로 실행합니다. float16 텐서가 절반의 메모리를 사용하므로 메모리 부족없이 배치 크기를 두 배로 늘릴 수 있습니다. 배치 크기를 늘리면 일반적으로 훈련 처리량, 즉 모델을 실행할 수 있는 초당 훈련 요소가 증가합니다.

GPU Tensor Cores 사용 보장

앞에서 언급했듯이 최신 NVIDIA GPU는 float16 행렬을 매우 빠르게 곱할 수 있는 Tensor Cores라는 특수 하드웨어 장치를 사용합니다. 그러나 Tensor Cores는 텐서의 특정 크기가 8의 배수 이어야 합니다. 다음 예에서는 텐서 코어를 사용하기 위해 8의 배수 이어야하는 경우에만 인수가 굵게 표시됩니다.

• tf.keras.layers.Dense(units=64)

• tf.keras.layers.Conv2d(filters=48, kernel_size=7, stride=3)

  • tf.keras.layers.Conv3d와 같은 다른 컨볼루셔널 레이어와 유사합니다

• tf.keras.layers.LSTM(units=64)

  • tf.keras.layers.GRU와 같은 다른 RNN과 유사합니다

• tf.keras.Model.fit(epochs=2, batch_size=128)

가능하면 Tensor Cores를 사용하십시오. 자세한 내용을 보려면 NVIDIA 딥 러닝 성능 안내서에 Tensor Cores 및 기타 Tensor Cores 관련 성능 정보를 사용하기위한 정확한 요구 사항이 설명되어 있습니다.

XLA

XLA는 혼합 정밀도 성능과 float32 성능을 어느 정도 더 향상시킬 수 있는 컴파일러입니다. 자세한 내용은 XLA 가이드를 참조하세요.

Cloud TPU 성능 팁

bfloat16 텐서는 메모리의 절반을 사용하므로 GPU와 마찬가지로 Cloud TPU를 사용할 때도 배치 크기를 두 배로 늘려야 합니다. 배치 크기를 두 배로 늘리면 학습 처리량이 증가할 수 있습니다.

TPU는 최적의 성능을 얻기 위해 다른 혼합 정밀도 조정이 필요하지 않습니다. TPU는 이미 XLA의 사용을 필요로 합니다. 특정한 크기를 $128$의 배수로 만들면 TPU 성능 이점을 얻을 수 있지만 혼합 정밀도의 경우에도 마찬가지이기 때문에 이러한 이점은 float32에도 동일하게 적용됩니다. float32와 혼합 정밀도에 모두 적용되는 일반적인 TPU 성능 팁은 Cloud TPU 성능 가이드를 참조하세요.

요약

• 계산 기능이 7.0 이상인 TPU, NVIDIA GPU 또는 AMX 명령어를 지원하는 인텔 CPU를 사용하는 경우 성능이 최대 3배 향상되므로 혼합 정밀도를 사용해야 합니다.

• 다음과 같이 혼합 정밀도를 사용할 수 있습니다.

  # On TPUs and CPUs, use 'mixed_bfloat16' instead
  mixed_precision.set_global_policy('mixed_float16')
  

• 모델이 softmax로 끝나는 경우 float32인지 확인합니다. 모델이 무엇으로 끝나는지에 관계없이 출력이 float32인지 확인합니다.

• mixed_float16으로 사용자 정의 훈련 루프를 사용하는 경우 위 코드 외에도 tf.keras.mixed_precision.LossScaleOptimizer로 옵티마이저를 래핑해야 합니다. 그런 다음 optimizer.get_scaled_loss를 호출하여 손실을 스케일링하고 optimizer.get_unscaled_gradients를 사용하여 그래디언트의 스케일링을 해제합니다.

• mixed_bfloat16과 함께 사용자 정의 훈련 루프를 사용하는 경우 위에서 언급한 global_policy를 설정하는 것으로도 충분합니다.

• 평가 정확성이 떨어지지 않으면 훈련 배치 크기를 두 배로 늘립니다

• GPU에서 성능을 최대화하려면 대부분의 텐서 차원이 $8$의 배수가 되도록 합니다.

tf.keras.mixed_precision API를 사용한 혼합 정밀도의 예는 훈련 성능과 관련된 함수 및 클래스. 자세한 내용은 Transformer와 같은 공식 모델을 확인하세요.