get dtype by global policy

google · Nov 20, 2020 · 4c83fbe · 4c83fbe
1 parent 902e157
commit 4c83fbe
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/efficientdet/dataloader.py b/efficientdet/dataloader.py
@@ -348,9 +348,8 @@ def dataset_parser(self, value, example_decoder, anchor_labeler, params):
       classes = pad_to_fixed_size(classes, -1,
                                   [self._max_instances_per_image, 1])
       if params['mixed_precision']:
-        precision = utils.get_precision(params['strategy'],
-                                        params['mixed_precision'])
-        dtype = precision.split('_')[-1]
+        dtype = (
+          tf.keras.mixed_precision.experimental.global_policy().compute_dtype)
         image = tf.cast(image, dtype=dtype)
         box_targets = tf.nest.map_structure(
             lambda box_target: tf.cast(box_target, dtype=dtype), box_targets)

diff --git a/efficientdet/keras/train_lib.py b/efficientdet/keras/train_lib.py
@@ -624,9 +624,7 @@ class and box losses from all levels.
     """
     # Sum all positives in a batch for normalization and avoid zero
     # num_positives_sum, which would lead to inf loss during training
-    precision = utils.get_precision(self.config.strategy,
-                                    self.config.mixed_precision)
-    dtype = precision.split('_')[-1]
+    dtype = tf.keras.mixed_precision.experimental.global_policy().compute_dtype
     num_positives_sum = tf.reduce_sum(labels['mean_num_positives']) + 1.0
     positives_momentum = self.config.positives_momentum or 0
     if positives_momentum > 0: