1# Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""Functional tests for deterministic BiasAdd.""" 16 17import numpy as np 18 19from absl.testing import parameterized 20 21from tensorflow.python.eager import backprop 22from tensorflow.python.eager import context 23from tensorflow.python.framework import config 24from tensorflow.python.framework import constant_op 25from tensorflow.python.framework import dtypes 26from tensorflow.python.framework import test_util 27from tensorflow.python.kernel_tests.nn_ops import bias_op_base 28from tensorflow.python.ops import array_ops 29from tensorflow.python.ops import gradients_impl 30from tensorflow.python.ops import nn_ops 31from tensorflow.python.platform import test 32 33 34class BiasAddDeterministicTest(bias_op_base.BiasAddTestBase, 35 parameterized.TestCase): 36 37 def _makeShapeTuple(self, batch_size, channel_count, data_rank, data_dim, 38 data_layout): 39 data_dims = data_rank * (data_dim,) 40 if data_layout == 'channels_first': 41 shape = (batch_size,) + (channel_count,) + data_dims 42 elif data_layout == 'channels_last': 43 shape = (batch_size,) + data_dims + (channel_count,) 44 else: 45 raise ValueError('Unknown data format') 46 return shape 47 48 def _dataFormatFromDataLayout(self, data_layout=None): 49 if data_layout == 'channels_first': 50 return 'NCHW' 51 elif data_layout == 'channels_last': 52 return 'NHWC' 53 else: 54 raise ValueError('Unknown data_layout') 55 56 def _randomNDArray(self, shape): 57 return 2 * np.random.random_sample(shape) - 1 58 59 def _randomDataOp(self, shape, data_type): 60 return constant_op.constant(self._randomNDArray(shape), dtype=data_type) 61 62 @parameterized.named_parameters( 63 *test_util.generate_combinations_with_testcase_name( 64 # With the selected layer configuration, at least in TensorFlow 65 # version 2.0, when data_layout='channels_last', bias_add operates 66 # deterministically by default. I don't know if this is true for 67 # all layer configurations. These cases are still being tested here, 68 # for completeness. 69 data_layout=['channels_first', 'channels_last'], 70 data_rank=[1, 2, 3], 71 data_type=[dtypes.float16, dtypes.float32, dtypes.float64])) 72 @test_util.run_in_graph_and_eager_modes 73 @test_util.run_cuda_only 74 def testDeterministicGradients(self, data_layout, data_rank, data_type): 75 with self.session(force_gpu=True): 76 # Using a cached_session with force_gpu=True does not work at the time 77 # of writing (2019-12-10). Before the @parameterized.named_parameters 78 # decorator was added, this non-cached session context was set outside 79 # the iteration loops for the parameter combinations, and so was re-used. 80 seed = ( 81 hash(data_layout) % 256 + hash(data_rank) % 256 + 82 hash(data_type) % 256) 83 np.random.seed(seed) 84 batch_size = 10 85 channel_count = 8 86 data_dim = 14 87 input_shape = self._makeShapeTuple(batch_size, channel_count, data_rank, 88 data_dim, data_layout) 89 bias_shape = (channel_count,) 90 output_shape = input_shape 91 input_val = self._randomDataOp(input_shape, data_type) 92 bias_val = self._randomDataOp(bias_shape, data_type) 93 data_format = self._dataFormatFromDataLayout(data_layout) 94 repeat_count = 5 95 if context.executing_eagerly(): 96 97 def bias_gradients(local_seed): 98 np.random.seed(local_seed) 99 upstream_gradients = self._randomDataOp(output_shape, data_type) 100 with backprop.GradientTape(persistent=True) as tape: 101 tape.watch(bias_val) 102 bias_add_output = nn_ops.bias_add( 103 input_val, bias_val, data_format=data_format) 104 gradient_injector_output = bias_add_output * upstream_gradients 105 return tape.gradient(gradient_injector_output, bias_val) 106 107 for i in range(repeat_count): 108 local_seed = seed + i # select different upstream gradients 109 result_a = bias_gradients(local_seed) 110 result_b = bias_gradients(local_seed) 111 self.assertAllEqual(result_a, result_b) 112 else: # graph mode 113 upstream_gradients = array_ops.placeholder( 114 data_type, shape=output_shape, name='upstream_gradients') 115 bias_add_output = nn_ops.bias_add( 116 input_val, bias_val, data_format=data_format) 117 gradient_injector_output = bias_add_output * upstream_gradients 118 # The gradient function behaves as if grad_ys is multiplied by the op 119 # gradient result, not passing the upstram gradients through the op's 120 # gradient generation graph. This is the reason for using the 121 # gradient injector 122 bias_gradients = gradients_impl.gradients( 123 gradient_injector_output, 124 bias_val, 125 grad_ys=None, 126 colocate_gradients_with_ops=True)[0] 127 for i in range(repeat_count): 128 feed_dict = {upstream_gradients: self._randomNDArray(output_shape)} 129 result_a = bias_gradients.eval(feed_dict=feed_dict) 130 result_b = bias_gradients.eval(feed_dict=feed_dict) 131 self.assertAllEqual(result_a, result_b) 132 133 # TODO(duncanriach): Re-enable the following three tests for the error checks 134 # after deterministic functionality is implemented at the CUDA kernel level. 135 def testInputDims(self): 136 pass 137 138 def testBiasVec(self): 139 pass 140 141 def testBiasInputsMatch(self): 142 pass 143 144 145if __name__ == '__main__': 146 # TODO(reedwm): Merge this file with bias_op_base.py and bias_op_test.py 147 config.enable_op_determinism() 148 test.main() 149