def testReluGrad(self, use_mlir): if use_mlir: SetTracingImplementation("mlir") def model(t): with tape_lib.GradientTape() as tape: tape.watch(t) result = unified_nn_ops.relu(t) grads = tape.gradient(result, t) return grads with context_lib.set_default(get_immediate_execution_context()): positive = TensorCastHelper(constant_op.constant([1.])) negative = TensorCastHelper(constant_op.constant([-1.])) model_fn = def_function.function(model) func_output = model_fn(positive) self.assertAllEqual(func_output.numpy(), [1.]) func_output = model_fn(negative) self.assertAllEqual(func_output.numpy(), [0.]) eager_output = model(positive) self.assertAllEqual(eager_output.numpy(), [1.]) eager_output = model(negative) self.assertAllEqual(eager_output.numpy(), [0.])
def __call__(self, *args, **kwargs): # Flatten arguments. flat_args = nest.flatten(args, expand_composites=True) flat_kwargs = nest.flatten(kwargs, expand_composites=True) all_args = flat_args + flat_kwargs # Trace outer_ctx = context_lib.get_default() ctx = NewTracingContext(self.name) with context_lib.set_default(ctx): # TODO(srbs): Iterating over list of inputs is a known performance # bottleneck. Add a pybind API for this. inputs = [ctx.AddParameter(arg.DataType()) for arg in all_args] structured_args = nest.pack_sequence_as(args, inputs[:len(flat_args)]) structured_kwargs = nest.pack_sequence_as(kwargs, inputs[len(flat_args):]) structured_outputs = self._python_func(*structured_args, **structured_kwargs) py_outputs = nest.flatten(structured_outputs, expand_composites=True) num_outputs = len(py_outputs) # TODO(srbs): Drop Nones before calling Finalize. finalized_f = ctx.Finalize(py_outputs) outer_ctx.RegisterFunction(finalized_f) # Build call op call_op = outer_ctx.CreateOperation(self.name, "") call_op.SetOpName(self.name) for arg in all_args: call_op.AddInput(arg) call_op_outputs = call_op.Execute(num_outputs) # Cleanup outer_ctx.RemoveFunction(self.name) return nest.pack_sequence_as(structured_outputs, call_op_outputs)
def testLog1p(self, use_mlir): if use_mlir: SetTracingImplementation("mlir") def model(a): return unified_math_ops.log1p(a) with context_lib.set_default(get_immediate_execution_context()): a = TensorCastHelper(constant_op.constant([1.])) func_output = def_function.function(model)(a) self.assertArrayNear(func_output.numpy(), [0.69314], 0.001) eager_output = model(a) self.assertArrayNear(eager_output.numpy(), [0.69314], 0.001)
def testNeg(self, use_mlir): if use_mlir: SetTracingImplementation("mlir") def model(a): return unified_math_ops.neg(a) with context_lib.set_default(get_immediate_execution_context()): a = TensorCastHelper(constant_op.constant([2.])) func_output = def_function.function(model)(a) self.assertAllEqual(func_output.numpy(), [-2.]) eager_output = model(a) self.assertAllEqual(eager_output.numpy(), [-2.])
def testDivNoNan(self, use_mlir): if use_mlir: SetTracingImplementation("mlir") def model(a, b): return unified_math_ops.div_no_nan(a, b) with context_lib.set_default(get_immediate_execution_context()): a = TensorCastHelper(constant_op.constant([2.])) b = TensorCastHelper(constant_op.constant([4.])) func_output = def_function.function(model)(a, b) self.assertArrayNear(func_output.numpy(), [0.5], 0.001) eager_output = model(a, b) self.assertArrayNear(eager_output.numpy(), [0.5], 0.001)
def testAdd(self, use_mlir): if use_mlir: SetTracingImplementation("mlir") def model(a, b): return math_ops.add(a, b) with context_lib.set_default(get_immediate_execution_context()): a = TensorCastHelper(constant_op.constant([1., 2.])) b = TensorCastHelper(constant_op.constant([3., 4.])) func_output = def_function.function(model)(a, b) self.assertAllEqual(func_output.numpy(), [4., 6.]) eager_output = model(a, b) self.assertAllEqual(eager_output.numpy(), [4., 6.])
def testAdd(self, use_tfrt, use_mlir): if use_mlir: SetTracingImplementation("mlir") def model(a, b): return math_ops.add(a, b) eager_ctx = NewImmediateExecutionContext(use_tfrt) with context_lib.set_default(eager_ctx): a = eager_ctx.CreateFloatScalarHandle(1.) b = eager_ctx.CreateFloatScalarHandle(2.) func_output = def_function.function(model)(a, b) self.assertAllEqual(func_output.numpy(), 3.0) eager_output = model(a, b) self.assertAllEqual(eager_output.numpy(), 3.0)
def _computeMnistMlpGrads(self, math_ops_lib, nn_ops_lib, backprop_lib, cast, num_iters, hidden_layers, hidden_size, batch_size): batch_size = 1 image_size = 28 * 28 num_classes = 10 def model(x, hidden_weights, softmax_weight, labels): with backprop_lib.GradientTape() as tape: for weight in hidden_weights + [softmax_weight]: tape.watch(weight) for hidden_weight in hidden_weights: x = math_ops_lib.mat_mul(x, hidden_weight) x = nn_ops_lib.relu(x) logits = math_ops_lib.mat_mul(x, softmax_weight) loss = nn_ops_lib.sparse_softmax_cross_entropy_with_logits( logits=logits, labels=labels) grads = tape.gradient(loss, hidden_weights + [softmax_weight]) return grads x = maybe_cast(array_ops.ones([batch_size, image_size]), cast) hidden_weights = [] for i in range(hidden_layers): hidden_weights.append( maybe_cast( random_ops.random_uniform( [hidden_size if i else image_size, hidden_size]), cast)) softmax_weight = maybe_cast( random_ops.random_uniform([hidden_size, num_classes]), cast) labels = maybe_cast(array_ops.zeros([batch_size], dtype=dtypes.int32), cast) with context_lib.set_default(get_immediate_execution_context()): # Warm up. for _ in range(10): model(x, hidden_weights, softmax_weight, labels) runtimes = timeit.repeat( lambda: model(x, hidden_weights, softmax_weight, labels), repeat=num_iters, number=10) return min(runtimes) / 10
def testLog1pGrad(self, use_mlir): if use_mlir: SetTracingImplementation("mlir") def model(a): with tape_lib.GradientTape() as tape: tape.watch(a) result = unified_math_ops.log1p(a) grads = tape.gradient(result, a) return grads with context_lib.set_default(get_immediate_execution_context()): a = TensorCastHelper(constant_op.constant([1.])) func_outputs = def_function.function(model)(a) self.assertArrayNear(func_outputs.numpy(), [0.5], 0.001) eager_outputs = model(a) self.assertArrayNear(eager_outputs.numpy(), [0.5], 0.001)
def testAddGrad(self, use_mlir): if use_mlir: SetTracingImplementation("mlir") def model(a, b): with tape_lib.GradientTape() as tape: tape.watch(a) tape.watch(b) result = unified_math_ops.add(a, b) grads = tape.gradient(result, [a, b]) return grads with context_lib.set_default(get_immediate_execution_context()): a = TensorCastHelper(constant_op.constant([1., 2.])) b = TensorCastHelper(constant_op.constant([3., 4.])) func_outputs = def_function.function(model)(a, b) self.assertAllEqual(func_outputs[0].numpy(), [1.0, 1.0]) self.assertAllEqual(func_outputs[1].numpy(), [1.0, 1.0]) eager_outputs = model(a, b) self.assertAllEqual(eager_outputs[0].numpy(), [1.0, 1.0]) self.assertAllEqual(eager_outputs[1].numpy(), [1.0, 1.0])
def testDivNoNanGrad(self, use_mlir): if use_mlir: SetTracingImplementation("mlir") def model(a, b): with tape_lib.GradientTape() as tape: tape.watch(a) tape.watch(b) result = unified_math_ops.div_no_nan(a, b) grads = tape.gradient(result, [a, b]) return grads with context_lib.set_default(get_immediate_execution_context()): a = TensorCastHelper(constant_op.constant([2.])) b = TensorCastHelper(constant_op.constant([4.])) func_outputs = def_function.function(model)(a, b) self.assertArrayNear(func_outputs[0].numpy(), [0.25], 0.001) self.assertArrayNear(func_outputs[1].numpy(), [-0.125], 0.001) eager_outputs = model(a, b) self.assertArrayNear(eager_outputs[0].numpy(), [0.25], 0.001) self.assertArrayNear(eager_outputs[1].numpy(), [-0.125], 0.001)
def testAddGrad(self, use_tfrt, use_mlir): if use_mlir: SetTracingImplementation("mlir") def model(a, b): with tape_lib.GradientTape() as tape: tape.watch(a) tape.watch(b) result = math_ops.add(a, b) grads = tape.gradient(result, [a, b]) return grads eager_ctx = NewImmediateExecutionContext(use_tfrt) with context_lib.set_default(eager_ctx): a = eager_ctx.CreateFloatScalarHandle(1.) b = eager_ctx.CreateFloatScalarHandle(2.) func_outputs = def_function.function(model)(a, b) self.assertAllEqual(func_outputs[0].numpy(), 1.0) self.assertAllEqual(func_outputs[1].numpy(), 1.0) eager_outputs = model(a, b) self.assertAllEqual(eager_outputs[0].numpy(), 1.0) self.assertAllEqual(eager_outputs[1].numpy(), 1.0)
def __enter__(self): """Enters a context inside which operations are recorded on this tape.""" self._ctx_manager = context_stack.set_default(self._tape_context) self._ctx_manager.__enter__() return self