Автоматическая смешанная точность от NVidia может быть способом до go.
Из того, что я собрал с 1.14, оно поддерживается (поддерживается) в апстриме. Все, что вам нужно сделать, это обернуть ваш оптимизатор следующим образом:
opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt)
Вам также может потребоваться установить специфику c environment variable в вашем скрипте Python, а именно:
os.environ[‘TF_ENABLE_AUTO_MIXED_PRECISION’] = ‘1’
Выше уже должны использоваться хорошие методы обучения смешанной точности (например, масштабирование потерь, ведение float32, где необходимо и т. Д. c.).
Хорошим ресурсом для этого решения должна быть официальная документация NVidia .
Собраны некоторые другие ресурсы, которые также могут быть полезны (хотя, похоже, не указывают, что вам придется делать что-то еще) здесь , здесь или здесь .
Я бы посоветовал не использовать ручное приведение, поскольку вы можете легко потерять точность (например, в статистике BatchNorm, используемой при выводе), если вы не знаете входы и выходы из спецификаций c слои.
Кроме того, вы также можете проверить тип Google bfloat16 (мозг-плавающий), который имеет exponent часть float32 (8 бит) и меньшую дробь. Это позволяет ему сохранять больший диапазон значений (например, при вычислении градиентов) по сравнению с float16, что позволяет избежать loss scaling.
Выше (bfloat16) должен быть полезен, в основном, в TPU, поддержка AFAIK NVidia GPU не слишком велика (кто-то поправит меня, если я ошибаюсь). Некоторая информация здесь .