Expected Data Types in Mixed Precision Cheatsheet

When training neural networks with the Keras API, we care about the data types and computation types since they are relevant to the convergence (numeric stability) and performance (memory footprint and computation efficiency). There are multiple “knobs” that we can turn to change the types by:

  1. Setting dtype of input tensors, or explicitly tf.cast the tensors;
  2. Setting dtype of the Keras Layer which defines data type of the layer’s computations and weights;
  3. Using the environment variable TF_FP16_CONV_USE_FP32_COMPUTE for the computation data type;
  4. Using the mixed precision policy globally: mixed_precision.set_global_policy('mixed_float16').

This may seem a bit confusing and it may not be clear how different settings affect each other and what we should expect about the actual weight/output data type and computation data type.

Therefore, I am trying to sweep through all possible combinations of the settings (fp16 or fp32) and the table below summarizes the obtained weight/output data type and computation data type for them and I hope the examples help.

Layer Input TF_FP16_CONV_USE _FP32_COMPUTE Weight Computation Output
fp32 fp32 1 fp32 fp32 fp32
fp32 fp32 0 fp32 fp32 fp32
fp32 fp16(⇒fp32) 1 fp32 fp32 fp32
fp32 fp16(⇒fp32) 0 fp32 fp32 fp32
fp16 fp32(⇒fp16) 1 fp16 fp32 fp16
fp16 fp32(⇒fp16) 0 fp16 fp16 fp16
fp16 fp16 1 fp16 fp32 fp16
fp16 fp16 0 fp16 fp16 fp16
mixed fp32(⇒fp16) 1 fp32(⇒fp16) fp32 fp16
mixed fp32(⇒fp16) 0 fp32(⇒fp16) fp16 fp16
mixed fp16 1 fp32(⇒fp16) fp32 fp16
mixed fp16 0 fp32(⇒fp16) fp16 fp16

A(=>B) means the data is stored in A but will be automatically cast to B.

Basically we can summarize the behavior with four rules:

  • The layer dtype determines the actual input/weight dtype.
  • “mixed” mainly refers to the weight that can be stored in fp32 but computed in fp16.
  • The env var determines the computation dtype when layer dtype is fp16/mixed.
  • The output dtype is fp32 only when layer dtype is fp32.

2022

CUDA Tips: nvcc’s -code, -arch, -gencode

1 minute read

Introduction People may feel confused by the options of -code, -arch, -gencode when compiling their CUDA codes. Although the official guidance explains the d...

Back to top ↑

2021

Demystifying the BatchNorm-Add-ReLU Fusion

2 minute read

Introduction My previous post, “Demystifying the Conv-Bias-ReLU Fusion”, has introduced a common fusion pattern in deep learning models. This post, on the ot...

Topological Sorting Explained

5 minute read

Introduction Recently I was working on a project related to the operation fusion in Tensorflow. My previous posts have covered several topics, such as how to...

Demystifying the Conv-Bias-ReLU Fusion

6 minute read

Introduction My previous post, “Fused Operations in Tensorflow”, introduced the basics of operation fusion in deep learning by showing how to enable the grap...

Fused Operations in Tensorflow

5 minute read

Introduction The computations in deep learning models are usually represented by a graph. Typically, operations in the graph are executed one by one, and eac...

Back to top ↑

2020

Inside Normalizations of Tensorflow

5 minute read

Introduction Recently I came across with optimizing the normalization layers in Tensorflow. Most online articles are talking about the mathematical definitio...

Back to top ↑