Does tensorflow use tensor cores?

[MiszS]

I'm making a CNN in python to classify images, and I'm wondering if tensorflow actually uses the tensor cores on my rtx 3090 for training, or does it just use the regular cuda cores?

 

[AnAverageName]

In short, as long as you're using TF 2.4+, yes. TF will automatically use the tensor cores on Ampere GPUs for supported supported FP32 operations using TensorFloat math mode.

 

You can check that by running:

import tensorflow as tf 
print(tf.config.experimental.tensor_float_32_execution_enabled())  

 

Ampere tensor cores also support FP16/BP16 operations so you may be able to get some more performance (or at least speed) when training by switching to mixed precision training where multiplications are done in 16-bit FP16 and accumulations in 32-bit FP32, halving the memory traffic while still using the Tensor Cores.

 

You can enable that with:

from tensorflow.keras import mixed_precision
mixed_precision.set_global_policy('mixed_float16')

 

End of the day some stuff will still run on CUDA though (any operations not supported by TF32 or with mixed precision kernels like FP32/FP64).

 

Hope this helps. Let me know if you have any questions!