Отсутствует 0-й вывод узла... При попытке использовать bfloat16 в tensorflow 2 ⇐ Python

[Anonymous]

Итак, я пытаюсь преобразовать существующий проект для использования bfloat16, поскольку это позволяет коду работать на тензорных ядрах. Я использую mix_precision.set_global_policy('mixed_bfloat16'), которого, согласно документации keras, достаточно и он должен работать. Однако при использовании этой строки я получаю следующую ошибку:

Код: Выделить всё

2 root error(s) found.
(0) INTERNAL:  Missing 0-th output from node model/sequential/conv2d_transpose/conv2d_transpose
(defined at C:\devlibs\Python\lib\site-packages\keras\backend.py:5530)

[[model/sequential/conv2d_transpose_4/BiasAdd/_76]]
(1) INTERNAL:  Missing 0-th output from node model/sequential/conv2d_transpose/conv2d_transpose
(defined at C:\devlibs\Python\lib\site-packages\keras\backend.py:5530)

0 successful operations.
0 derived errors ignored. [Op:__inference_predict_function_1627]

Errors may have originated from an input operation.
Input Source operations connected to node model/sequential/conv2d_transpose/conv2d_transpose:
In[0] model/sequential/conv2d_transpose/stack (defined at C:\devlibs\Python\lib\site-packages\keras\layers\convolutional.py:1333)
In[1] model/sequential/conv2d_transpose/conv2d_transpose/Cast (defined at C:\devlibs\Python\lib\site-packages\keras\mixed_precision\autocast_variable.py:146)
In[2] model/sequential/reshape/Reshape (defined at C:\devlibs\Python\lib\site-packages\keras\layers\core\reshape.py:126)

Эта ошибка возникает на первом уровне во время обучения, но если я заставлю этот слой использовать float32, ошибка перейдет на следующий уровень (по сути, все слои несовместимы с bfloat16). Возможно, стоит отметить, что мои входные данные находятся в формате uint8.

Подробнее здесь: https://stackoverflow.com/questions/700 ... nsorflow-2