Commit 48b60914 authored by Vinh Nguyen's avatar Vinh Nguyen
Browse files

minor fix automatic mixed precision - resnet

parent 80af2a7b
......@@ -32,7 +32,6 @@ from official.vision.image_classification import common
from official.vision.image_classification import imagenet_preprocessing
from official.vision.image_classification import resnet_model
from official.vision.image_classification import trivial_model
import pdb
LR_SCHEDULE = [ # (multiplier, epoch to start) tuples
(1.0, 5), (0.1, 30), (0.01, 60), (0.001, 80)
......@@ -185,11 +184,11 @@ def run(flags_obj):
optimizer = tf.keras.mixed_precision.experimental.LossScaleOptimizer(
optimizer, loss_scale=flags_core.get_loss_scale(flags_obj,
default_for_fp16=128))
pdb.set_trace()
if flags_obj.fp16_implementation == "graph_rewrite":
# Note: when flags_obj["fp16_implementation"] == "graph_rewrite",
# Note: when flags_obj.fp16_implementation == "graph_rewrite",
# dtype as determined by flags_core.get_tf_dtype(flags_obj) would be 'float32'
#
# which will ensure tf.keras.mixed_precision and tf.train.experimental.enable_mixed_precision_graph_rewrite
# does not double up.
optimizer = tf.train.experimental.enable_mixed_precision_graph_rewrite(optimizer)
if flags_obj.use_trivial_model:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment