How to support mixed precision in custom Tensorflow layers?

user_1037200 0
Viewed 778

When developing my own custom layers for tf.keras: how am I supposed to support mixed precision?

The documentation of mixed precision - a feature which is currently marked as experimental in Tensorflow 2.2 - only explains how to use it from a consumers perspective with predefined layers such as the tf.keras.layers.Dense one.

I already tried to guess it myself and found two - maybe relevant - details:

  • The dtype property stays as float32 by default when using 16-bit mixed precision.

  • There is a mixed_precision.get_layer_policy(layer) method (see docs) and a mixed_precision.global_policy() method (see docs) which could be used to retrieve the configured compute_dtype and variable_dtype.

Am I supposed to use the above get_layer_policy-method and just cast my variables into compute_dtype within the call(...) method of my layer? (And pass variable_dtype in my layers build(...) method to add_weight(...) when creating variables?)

For example, here is naive sample implementation of a standard dense neuron layer:

  def call(self, input):
    policy = mixed_precision.get_layer_policy(self)
    bias = tf.cast(self._bias, policy.compute_dtype)
    weights = tf.cast(self._weights, policy.compute_dtype)
    y = tf.nn.bias_add(tf.matmul(input, weights), bias)
    outputs = self._activation(y)
    return outputs

Sure, nobody would implement such basic stuff themselves, that one is just for demonstration. But, would this be the way the Tensorflow team expects us to implement the call(...) methods of our custom layers?

tensorflow2.x tf.keras tensorflow python
1 Answers

This guide (slide 13-14) from nvidia mentions custom layers for mixed precision training.

You have to implement the method cast_input(). In this example the layer is casted to float16 when mixed precision is enabled:

class CustomBiasLayer(tf.keras.layers.Layer):

 def build(self, _):
   self.v = self.add_weight('v', ())
   self.built = True
  
 def call(self, inputs):
   return inputs + self.v

 def cast_inputs(self, inputs):
   # Casts to float16, the policy's lowest-precision dtype
   return self._mixed_precision_policy.cast_to_lowest(inputs)

I have not tried this myself, so please let me know if this works for you.

user_116406950
Related