def testVariableName(self): """Test if custom name for the variable scope is propagated correctly.""" g = ops.Graph() with g.as_default(): a = variables.Variable(2.2, name='var_a') b = variables.Variable(3.3, name='var_b') d = variables.Variable(4.4, name='var_b') with g.name_scope('scope1'): with bfloat16.bfloat16_scope('bf16'): a = math_ops.cast(a, dtypes.bfloat16) b = math_ops.cast(b, dtypes.bfloat16) c = math_ops.add(a, b, name='addition') with bfloat16.bfloat16_scope(): d = math_ops.cast(d, dtypes.bfloat16) math_ops.add(c, d, name='addition') g_ops = g.get_operations() ops_name = [] for op in g_ops: ops_name.append(str(op.name)) self.assertIn('scope1/bf16/addition', ops_name) self.assertIn('scope1/bf16/Cast', ops_name) self.assertIn('scope1/addition', ops_name) self.assertIn('scope1/Cast', ops_name)
def tpu_subgraph_predict(): if use_bfloat16: with bfloat16_scope(): return tpu.rewrite(tpu_subgraph_predict_fn, [preprocessed_inputs, true_image_shapes]) else: return tpu.rewrite(tpu_subgraph_predict_fn, [preprocessed_inputs, true_image_shapes])
def testRequestedDType(self): """Test if requested dtype is honored in the getter. """ with bfloat16.bfloat16_scope() as scope: v1 = variable_scope.get_variable("v1", []) self.assertEqual(v1.dtype.base_dtype, dtypes.float32) v2 = variable_scope.get_variable("v2", [], dtype=dtypes.bfloat16) self.assertEqual(v2.dtype.base_dtype, dtypes.bfloat16) self.assertEqual([dtypes.float32, dtypes.float32], [v.dtype.base_dtype for v in scope.global_variables()])
def testRequestedDType(self): """Test if requested dtype is honored in the getter. """ with bfloat16.bfloat16_scope() as scope: v1 = variable_scope.get_variable("v1", []) self.assertEqual(v1.dtype.base_dtype, dtypes.float32) v2 = variable_scope.get_variable("v2", [], dtype=dtypes.bfloat16) self.assertEqual(v2.dtype.base_dtype, dtypes.bfloat16) self.assertEqual([dtypes.float32, dtypes.float32], [v.dtype.base_dtype for v in scope.global_variables()])
def predict_tpu_subgraph(preprocessed_inputs, true_image_shapes): """Wraps over the CPU version of `predict()`. This builds a same graph as the original `predict()`, manipulates result tensors' dimensions to be memory efficient on TPU, and returns them as list of tensors. Args: preprocessed_inputs: A 4D tensor of shape (batch, channels, height, width) true_image_shapes: True image shapes tensor. Returns: A Python list of tensors: box_encodings: 3D tensor of shape (code_size, batch_size, num_anchors) class_predictions_with_background: 3D tensor, shape (num_classes + 1, batch_size, num_anchors) anchors: 2D tensor of shape (4, num_anchors) """ # Dimshuffle: (b, c, h, w) -> (b, h, w, c) preprocessed_inputs = tf.transpose(preprocessed_inputs, perm=[0, 2, 3, 1]) if use_bfloat16: with bfloat16_scope(): prediction_dict = detection_model.predict( preprocessed_inputs, true_image_shapes) else: prediction_dict = detection_model.predict(preprocessed_inputs, true_image_shapes) # Dimshuffle: (batch, anchors, depth) -> (depth, batch, anchors) return [ tf.transpose(prediction_dict[BOX_ENCODINGS], perm=[2, 0, 1]), tf.transpose(prediction_dict[CLASS_PREDICTIONS_WITH_BACKGROUND], perm=[2, 0, 1]), tf.transpose(prediction_dict[ANCHORS], perm=[1, 0]), ]
def testScopeName(self): """Test if name for the variable scope is propagated correctly.""" with bfloat16.bfloat16_scope() as bf: self.assertEqual(bf.name, "")
def testCustomScopeName(self): """Test if custom name for the variable scope is propagated correctly.""" name = 'bfloat16' with bfloat16.bfloat16_scope('bfloat16') as bf: self.assertEqual(bf.name, name)
def testScopeName(self): """Test if name for the variable scope is propogated correctly. """ with bfloat16.bfloat16_scope() as bf: self.assertEqual(bf.name, "")