Смешанная точность

Посмотреть на TensorFlow.org Запустить в Google Colab Посмотреть исходный код на GitHub Скачать блокнот

Обзор

Смешанная точность — это использование 16-битных и 32-битных типов с плавающей запятой в модели во время обучения, чтобы она работала быстрее и использовала меньше памяти. Сохраняя определенные части модели в 32-битных типах для числовой стабильности, модель будет иметь меньшее время шага и одинаково хорошо обучаться с точки зрения показателей оценки, таких как точность. В этом руководстве описывается, как использовать API смешанной точности Keras для ускорения ваших моделей. Использование этого API может повысить производительность более чем в 3 раза на современных GPU и на 60% на TPU.

Сегодня в большинстве моделей используется dtype float32, занимающий 32 бита памяти. Однако есть два типа d более низкой точности, float16 и bfloat16, каждый из которых вместо этого занимает 16 бит памяти. Современные ускорители могут быстрее выполнять операции в 16-битных dtypes, поскольку они имеют специализированное оборудование для выполнения 16-битных вычислений, а 16-битные dtypes можно быстрее считывать из памяти.

Графические процессоры NVIDIA могут выполнять операции в float16 быстрее, чем в float32, а TPU могут выполнять операции в bfloat16 быстрее, чем в float32. Таким образом, эти dtypes с более низкой точностью следует использовать, когда это возможно, на этих устройствах. Тем не менее, переменные и несколько вычислений все еще должны быть в float32 по числовым причинам, чтобы модель обучалась с тем же качеством. API смешанной точности Keras позволяет использовать сочетание float16 или bfloat16 с float32, чтобы получить преимущества производительности от float16/bfloat16 и преимущества числовой стабильности от float32.

Настраивать

import tensorflow as tf

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

Поддерживаемое оборудование

Хотя смешанная точность будет работать на большинстве аппаратных средств, она ускорит модели только на новейших графических процессорах NVIDIA и облачных TPU. Графические процессоры NVIDIA поддерживают сочетание float16 и float32, а TPU — сочетание bfloat16 и float32.

Среди графических процессоров NVIDIA те, у которых вычислительная мощность 7.0 или выше, получат наибольшую выгоду от смешанной точности, поскольку они имеют специальные аппаратные блоки, называемые тензорными ядрами, для ускорения умножения и свертки матриц с плавающей запятой16. Старые графические процессоры не дают преимущества в математической производительности при использовании смешанной точности, однако экономия памяти и пропускной способности может обеспечить некоторое ускорение. Вы можете посмотреть вычислительные возможности вашего GPU на веб-странице NVIDIA CUDA GPU . Примеры графических процессоров, которые больше всего выиграют от смешанной точности, включают графические процессоры RTX, V100 и A100.

Вы можете проверить свой тип графического процессора следующим образом. Команда существует только в том случае, если установлены драйверы NVIDIA, поэтому в противном случае следующее вызовет ошибку.

nvidia-smi -L
GPU 0: Tesla V100-SXM2-16GB (UUID: GPU-99e10c4d-de77-42ee-4524-6c41c4e5e47d)

Все облачные TPU поддерживают bfloat16.

Даже на процессорах и старых графических процессорах, где не ожидается ускорения, API-интерфейсы смешанной точности можно использовать для модульного тестирования, отладки или просто для опробования API. Однако на процессорах смешанная точность будет работать значительно медленнее.

Настройка политики dtype

Чтобы использовать смешанную точность в Keras, вам необходимо создать tf.keras.mixed_precision.Policy , обычно называемую политикой dtype . Политики Dtype определяют, в каких слоях dtypes будут работать. В этом руководстве вы создадите политику из строки 'mixed_float16' и установите ее в качестве глобальной политики. Это приведет к тому, что впоследствии созданные слои будут использовать смешанную точность со смесью float16 и float32.

policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)
INFO:tensorflow:Mixed precision compatibility check (mixed_float16): OK
Your GPU will likely run quickly with dtype policy mixed_float16 as it has compute capability of at least 7.0. Your GPU: Tesla V100-SXM2-16GB, compute capability 7.0

Короче говоря, вы можете напрямую передать строку в set_global_policy , что обычно и делается на практике.

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

Политика определяет два важных аспекта слоя: dtype, в котором выполняются вычисления слоя, и dtype переменных слоя. Выше вы создали политику mixed_float16 (т. е. mixed_precision.Policy , созданную путем передачи строки 'mixed_float16' в ее конструктор). С этой политикой слои используют вычисления float16 и переменные float32. Вычисления выполняются в формате float16 для повышения производительности, но переменные должны храниться в формате float32 для стабильности числовых значений. Вы можете напрямую запросить эти свойства политики.

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

Как упоминалось ранее, mixed_float16 значительно улучшит производительность графических процессоров NVIDIA с вычислительной мощностью не ниже 7.0. Политика будет работать на других GPU и CPU, но может не повысить производительность. Вместо этого для mixed_bfloat16 следует использовать политикуmixed_bfloat16.

Построение модели

Далее приступим к построению простой модели. Очень маленькие игрушечные модели обычно не выигрывают от смешанной точности, поскольку накладные расходы среды выполнения TensorFlow обычно доминируют над временем выполнения, что делает любое улучшение производительности на графическом процессоре незначительным. Поэтому давайте построим два больших Dense слоя по 4096 единиц каждый, если используется GPU.

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)
The model will run with 4096 units on a GPU

Каждый уровень имеет политику и по умолчанию использует глобальную политику. Таким образом, каждый из слоев Dense имеет политику mixed_float16 , поскольку ранее вы установили глобальную политику mixed_float16 . Это заставит плотные слои выполнять вычисления с плавающей запятой 16 и иметь переменные с плавающей точкой 32. Они приводят свои входные данные к типу с плавающей запятой16, чтобы выполнять вычисления с плавающей точкой16, что в результате приводит к тому, что их выходные данные имеют значение с плавающей запятой16. Их переменные имеют тип float32 и будут преобразованы в float16 при вызове слоев, чтобы избежать ошибок из-за несоответствия dtype.

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)
<Policy "mixed_float16">
x.dtype: float16
dense1.kernel.dtype: float32

Затем создайте выходные прогнозы. Обычно вы можете создавать выходные прогнозы следующим образом, но это не всегда численно стабильно с 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)
Outputs dtype: float16

Активация softmax в конце модели должна быть float32. Поскольку политика mixed_float16 —mixed_float16 , активация softmax обычно имеет dtype вычислений float16 и выводит тензоры 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)
Outputs dtype: float32

Передача dtype='float32' в конструктор слоя softmax переопределяет политику dtype слоя на политику float32 , которая выполняет вычисления и хранит переменные в float32. Эквивалентно, вместо этого вы могли бы передать dtype=mixed_precision.Policy('float32') ; слои всегда преобразовывают аргумент dtype в политику. Поскольку на уровне Activation нет переменных, переменная dtype политики игнорируется, но вычисление dtype политики, равное float32, приводит к тому, что softmax и выходные данные модели будут float32.

Добавление softmax float16 в середине модели — это нормально, но softmax в конце модели должен быть в float32. Причина в том, что если промежуточный тензор, идущий от softmax к убытку, равен float16 или bfloat16, могут возникнуть числовые проблемы.

Вы можете изменить dtype любого слоя на float32, передав dtype='float32' если вы считаете, что он не будет численно стабильным с вычислениями float16. Но обычно это необходимо только на последнем слое модели, так как большинство слоев имеют достаточную точность с 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. Вы не выполняете приведение к float16, поскольку деление на 255 выполняется процессором, который выполняет операции с плавающей запятой медленнее, чем операции с плавающей запятой32. В этом случае разница в производительности незначительна, но в целом вы должны запускать математику обработки ввода в float32, если она выполняется на процессоре. Первый уровень модели будет приводить входные данные к типу 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])
Epoch 1/5
6/6 [==============================] - 2s 78ms/step - loss: 4.9609 - accuracy: 0.4132 - val_loss: 0.6643 - val_accuracy: 0.8437
Epoch 2/5
6/6 [==============================] - 0s 34ms/step - loss: 0.7752 - accuracy: 0.7789 - val_loss: 0.3098 - val_accuracy: 0.9175
Epoch 3/5
6/6 [==============================] - 0s 34ms/step - loss: 0.3620 - accuracy: 0.8848 - val_loss: 0.3149 - val_accuracy: 0.8969
Epoch 4/5
6/6 [==============================] - 0s 34ms/step - loss: 0.2998 - accuracy: 0.9066 - val_loss: 0.2988 - val_accuracy: 0.9068
Epoch 5/5
6/6 [==============================] - 0s 33ms/step - loss: 0.2298 - accuracy: 0.9285 - val_loss: 0.5062 - val_accuracy: 0.8414
313/313 - 0s - loss: 0.5163 - accuracy: 0.8392
Test loss: 0.5163048505783081
Test accuracy: 0.8392000198364258

Обратите внимание, что модель печатает время на шаг в журналах: например, «25 мс/шаг». Первая эпоха может быть медленнее, так как TensorFlow тратит некоторое время на оптимизацию модели, но впоследствии время на шаг должно стабилизироваться.

Если вы запускаете это руководство в Colab, вы можете сравнить производительность смешанной точности с float32. Для этого измените политику mixed_float16 на float32 в разделе «Настройка политики dtype», а затем перезапустите все ячейки до этого момента. На графических процессорах с вычислительными возможностями 7.X вы должны увидеть значительное увеличение времени на шаг, что указывает на то, что смешанная точность ускорила модель. Не забудьте снова изменить политику на mixed_float16 и повторно запустить ячейки, прежде чем продолжить работу с руководством.

На графических процессорах с вычислительной мощностью не менее 8.0 (графические процессоры Ampere и выше) вы, скорее всего, не увидите улучшения производительности игрушечной модели в этом руководстве при использовании смешанной точности по сравнению с float32. Это связано с использованием TensorFloat-32 , который автоматически использует математику более низкой точности в некоторых операциях float32, таких как tf.linalg.matmul . TensorFloat-32 дает некоторые преимущества производительности смешанной точности при использовании float32. Однако в реальных моделях вы по-прежнему обычно будете наблюдать значительный прирост производительности за счет смешанной точности из-за экономии пропускной способности памяти и операций, которые TensorFloat-32 не поддерживает.

При использовании смешанной точности на TPU вы не увидите такого большого прироста производительности по сравнению со смешанной точностью на графических процессорах, особенно на графических процессорах до Ampere. Это связано с тем, что TPU выполняют определенные операции в bfloat16 под капотом даже с политикой dtype по умолчанию float32. Это похоже на то, как графические процессоры Ampere по умолчанию используют TensorFloat-32. По сравнению с графическими процессорами Ampere, TPU обычно демонстрируют меньший прирост производительности при смешанной точности на реальных моделях.

Для многих реальных моделей смешанная точность также позволяет удвоить размер пакета без нехватки памяти, поскольку тензоры float16 занимают половину памяти. Однако это не относится к этой игрушечной модели, поскольку вы, вероятно, можете запустить модель в любом dtype, где каждый пакет состоит из всего набора данных MNIST из 60 000 изображений.

Масштабирование потерь

Масштабирование убытков — это метод, который tf.keras.Model.fit автоматически выполняет с mixed_float16 , чтобы избежать числового недополнения. В этом разделе описывается, что такое масштабирование потерь, а в следующем разделе описывается, как использовать его с пользовательским циклом обучения.

Недостаточное и переполнение

Тип данных float16 имеет узкий динамический диапазон по сравнению с float32. Это означает, что значения выше \(65504\) будут переполняться до бесконечности, а значения ниже \(6.0 \times 10^{-8}\) будут уменьшаться до нуля. float32 и bfloat16 имеют гораздо более широкий динамический диапазон, так что переполнение и потеря значимости не являются проблемой.

Например:

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

На практике переполнение с float16 происходит редко. Кроме того, во время прямого прохода также редко возникает недолив. Однако при обратном проходе градиенты могут опускаться до нуля. Масштабирование убытков - это метод предотвращения этого недорасхода.

Обзор масштабирования потерь

Основная концепция масштабирования убытков проста: просто умножьте убыток на некоторое большое число, скажем, \(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

Выбор масштаба потерь может быть сложным. Если шкала потерь слишком мала, градиенты все равно могут опускаться до нуля. Если слишком высокое, возникает обратная проблема: градиенты могут выйти за пределы бесконечности.

Чтобы решить эту проблему, TensorFlow динамически определяет масштаб потерь, поэтому вам не нужно выбирать его вручную. Если вы используете tf.keras.Model.fit , масштабирование потерь выполняется за вас, поэтому вам не нужно выполнять какую-либо дополнительную работу. Если вы используете пользовательский цикл обучения, вы должны явно использовать специальную оболочку оптимизатора tf.keras.mixed_precision.LossScaleOptimizer , чтобы использовать масштабирование потерь. Это описано в следующем разделе.

Обучение модели с помощью пользовательского цикла обучения

До сих пор вы обучали модель Keras со смешанной точностью, используя tf.keras.Model.fit . Далее вы будете использовать смешанную точность с пользовательским циклом обучения. Если вы еще не знаете, что такое настраиваемый цикл обучения, сначала прочтите Руководство по настраиваемому обучению .

Запуск пользовательского цикла обучения со смешанной точностью требует двух изменений по сравнению с его запуском в float32:

  1. Постройте модель со смешанной точностью (вы уже сделали это)
  2. Явно используйте масштабирование потерь, mixed_float16 используетсяmixed_float16.

Для шага (2) вы будете использовать класс tf.keras.mixed_precision.LossScaleOptimizer , который обертывает оптимизатор и применяет масштабирование потерь. По умолчанию он динамически определяет масштаб потерь, поэтому вам не нужно его выбирать. LossScaleOptimizer следующим образом.

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

При желании можно выбрать явную шкалу потерь или иным образом настроить поведение масштабирования потерь, но настоятельно рекомендуется сохранить поведение масштабирования потерь по умолчанию, так как было обнаружено, что оно хорошо работает на всех известных моделях. См. документацию tf.keras.mixed_precision.LossScaleOptimizer , если вы хотите настроить поведение масштабирования потерь.

Затем определите объект потерь и tf.data.Dataset s:

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()))
Epoch 0: loss=4.869325160980225, test accuracy=0.7221999764442444
Epoch 1: loss=0.4893573224544525, test accuracy=0.878000020980835
Epoch 2: loss=0.36011582612991333, test accuracy=0.9440000057220459
Epoch 3: loss=0.27391332387924194, test accuracy=0.9318000078201294
Epoch 4: loss=0.247697651386261, test accuracy=0.933899998664856

Советы по производительности графического процессора

Вот несколько советов по повышению производительности при использовании смешанной точности на графических процессорах.

Увеличение размера партии

Если это не влияет на качество модели, попробуйте запустить с удвоенным размером пакета при использовании смешанной точности. Поскольку тензоры float16 используют половину памяти, это часто позволяет вам удвоить размер пакета без нехватки памяти. Увеличение размера пакета обычно увеличивает пропускную способность обучения, т. е. количество обучающих элементов в секунду, на которых может работать ваша модель.

Обеспечение использования тензорных ядер графического процессора

Как упоминалось ранее, современные графические процессоры NVIDIA используют специальный аппаратный блок под названием Tensor Cores, который может очень быстро умножать матрицы с плавающей запятой16. Однако тензорные ядра требуют, чтобы определенные размеры тензоров были кратны 8. В приведенных ниже примерах аргумент выделен жирным шрифтом тогда и только тогда, когда для использования тензорных ядер он должен быть кратен 8.

  • tf.keras.layers.Dense ( единицы = 64 )
  • tf.keras.layers.Conv2d ( фильтры = 48 , размер ядра = 7, шаг = 3)
    • И аналогично для других сверточных слоев, таких как tf.keras.layers.Conv3d
  • tf.keras.layers.LSTM ( единицы = 64 )
    • И аналогично для других RNN, таких как tf.keras.layers.GRU.
  • tf.keras.Model.fit (эпохи = 2, batch_size = 128 )

Вы должны попытаться использовать тензорные ядра, когда это возможно. Если вы хотите узнать больше, руководство NVIDIA по производительности глубокого обучения описывает точные требования для использования тензорных ядер, а также другую информацию о производительности, связанную с тензорными ядрами.

XLA

XLA — это компилятор, который может дополнительно повысить производительность смешанной точности, а также в меньшей степени производительность float32. Подробную информацию см. в руководстве XLA .

Советы по производительности облачного TPU

Как и в случае с графическими процессорами, вам следует попробовать удвоить размер пакета при использовании облачных TPU, поскольку тензоры bfloat16 используют половину памяти. Удвоение размера пакета может увеличить производительность обучения.

TPU не требуют какой-либо другой настройки смешанной точности для достижения оптимальной производительности. Они уже требуют использования XLA. TPU выигрывают от того, что некоторые размеры кратны \(128\), но это в равной степени относится к типу float32, как и к смешанной точности. Ознакомьтесь с руководством по производительности Cloud TPU , чтобы узнать общие советы по производительности TPU, которые применимы как к тензорам со смешанной точностью, так и к тензорам с плавающей запятой.

Резюме

  • Вам следует использовать смешанную точность, если вы используете TPU или графические процессоры NVIDIA с вычислительными возможностями не ниже 7.0, так как это повысит производительность до 3 раз.
  • Вы можете использовать смешанную точность со следующими строками:

    # On TPUs, 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 , чтобы отменить масштабирование градиентов.

  • Удвойте размер обучающей партии, если это не снижает точность оценки.

  • На графических процессорах убедитесь, что размеры большинства тензоров кратны \(8\) , чтобы максимизировать производительность.

Дополнительные примеры смешанной точности с использованием API tf.keras.mixed_precision можно найти в официальном репозитории моделей . Большинство официальных моделей, таких как ResNet и Transformer , будут работать со смешанной точностью при передаче --dtype=fp16 .