diff --git a/tensorflow/python/eager/benchmarks/resnet50/hvp_test.py b/tensorflow/python/eager/benchmarks/resnet50/hvp_test.py index c5282174ffa..fa86a98371d 100644 --- a/tensorflow/python/eager/benchmarks/resnet50/hvp_test.py +++ b/tensorflow/python/eager/benchmarks/resnet50/hvp_test.py @@ -35,7 +35,7 @@ def _forward_over_back_hvp(model, images, labels, vector): model.trainable_variables, vector) as acc: with tf.GradientTape() as grad_tape: logits = model(images, training=True) - loss = tf.losses.softmax_cross_entropy( + loss = tf.compat.v1.losses.softmax_cross_entropy( logits=logits, onehot_labels=labels) grads = grad_tape.gradient(loss, model.trainable_variables) return acc.jvp(grads) @@ -47,7 +47,7 @@ def _back_over_forward_hvp(model, images, labels, vector): with forwardprop.ForwardAccumulator( model.trainable_variables, vector) as acc: logits = model(images, training=True) - loss = tf.losses.softmax_cross_entropy( + loss = tf.compat.v1.losses.softmax_cross_entropy( logits=logits, onehot_labels=labels) return grad_tape.gradient(acc.jvp(loss), model.trainable_variables) @@ -55,7 +55,7 @@ def _back_over_forward_hvp(model, images, labels, vector): def _tf_gradients_forward_over_back_hvp(model, images, labels, vector): with tf.GradientTape() as grad_tape: logits = model(images, training=True) - loss = tf.losses.softmax_cross_entropy( + loss = tf.compat.v1.losses.softmax_cross_entropy( logits=logits, onehot_labels=labels) variables = model.trainable_variables grads = grad_tape.gradient(loss, variables) @@ -68,7 +68,7 @@ def _back_over_back_hvp(model, images, labels, vector): with tf.GradientTape() as outer_tape: with tf.GradientTape() as inner_tape: logits = model(images, training=True) - loss = tf.losses.softmax_cross_entropy( + loss = tf.compat.v1.losses.softmax_cross_entropy( logits=logits, onehot_labels=labels) grads = inner_tape.gradient(loss, model.trainable_variables) return outer_tape.gradient( diff --git a/tensorflow/python/eager/benchmarks/resnet50/resnet50_graph_test.py b/tensorflow/python/eager/benchmarks/resnet50/resnet50_graph_test.py index 2520df70b44..958756d418b 100644 --- a/tensorflow/python/eager/benchmarks/resnet50/resnet50_graph_test.py +++ b/tensorflow/python/eager/benchmarks/resnet50/resnet50_graph_test.py @@ -109,7 +109,7 @@ class ResNet50Benchmarks(tf.test.Benchmark): model = resnet50.ResNet50(data_format()) logits = model(images, training=True) - loss = tf.losses.softmax_cross_entropy( + loss = tf.compat.v1.losses.softmax_cross_entropy( logits=logits, onehot_labels=labels) optimizer = tf.train.GradientDescentOptimizer(learning_rate=1.0) train_op = optimizer.minimize(loss) diff --git a/tensorflow/python/eager/benchmarks/resnet50/resnet50_test.py b/tensorflow/python/eager/benchmarks/resnet50/resnet50_test.py index 8be6ab89766..6e8d7b61669 100644 --- a/tensorflow/python/eager/benchmarks/resnet50/resnet50_test.py +++ b/tensorflow/python/eager/benchmarks/resnet50/resnet50_test.py @@ -36,7 +36,7 @@ from tensorflow.python.eager.benchmarks.resnet50 import resnet50_test_util def compute_gradients(model, images, labels, num_replicas=1): with tf.GradientTape() as grad_tape: logits = model(images, training=True) - loss = tf.losses.softmax_cross_entropy( + loss = tf.compat.v1.losses.softmax_cross_entropy( logits=logits, onehot_labels=labels) tf.compat.v2.summary.write('loss', loss) if num_replicas != 1: diff --git a/tensorflow/python/eager/benchmarks/resnet50/resnet50_test_util.py b/tensorflow/python/eager/benchmarks/resnet50/resnet50_test_util.py index dacc6fb6a85..8d3b95053b5 100644 --- a/tensorflow/python/eager/benchmarks/resnet50/resnet50_test_util.py +++ b/tensorflow/python/eager/benchmarks/resnet50/resnet50_test_util.py @@ -35,8 +35,8 @@ def random_batch(batch_size, data_format): shape = (batch_size,) + shape num_classes = 1000 - images = tf.random_uniform(shape) - labels = tf.random_uniform( + images = tf.random.uniform(shape) + labels = tf.random.uniform( [batch_size], minval=0, maxval=num_classes, dtype=tf.int32) one_hot = tf.one_hot(labels, num_classes)