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"""Tests for deterministic cuDNN functionality.""" 16 17import collections 18 19import numpy as np 20 21from tensorflow.python.eager import backprop 22from tensorflow.python.framework import constant_op 23from tensorflow.python.framework import dtypes 24from tensorflow.python.framework import test_util 25from tensorflow.python.ops import nn_ops 26from tensorflow.python.platform import test 27 28# Notes: 29# 30# TensorFlow makes cuDNN run deterministically when op determinism is enabled 31# via tf.config.experimental.enable_op_determinism(). Additionally, setting the 32# environmental variable TF_CUDNN_DETERMINISTIC to 'true' or '1' makes cuDNN run 33# deterministically, although this environemtnal variable is deprecated and will 34# be removed in a future TensorFlow version. Unlike the enable_op_determinism() 35# function, the environmental variable only makes ops using cuDNN deterministic, 36# not all TensorFlow ops. 37# 38# Where both deterministic and non-deterministic cuDNN algorithms are available, 39# selecting determinitic operation will lead to only the deterministic 40# algorithms being chosen. Additionally, selecting deterministic operation will 41# result in a deterministic, or reproducible, selection of algorithms (for any 42# given layer configuration) for each of the forward and the two backward paths. 43# 44# These tests intend to confirm that deterministic algorithms are chosen (for 45# the back-prop paths) when desterministic operation is selected. The tested 46# configurations were first confirmed to produce non-deterministic results when 47# the above-mentioned environment variables are not set. 48# 49# Even though selecting determinitic operation should ensure that the same 50# algorithms, for a given layer configuration, are always used (i.e. that 51# algorithm selection is deterministic / reproducible), this is not tested. 52 53# TODO(duncanriach): Add test for deterministic cuDNN max-pooling 54 55LayerShapeNHWC = collections.namedtuple('LayerShapeNHWC', 56 'batch, height, width, channels') 57FilterShape2D = collections.namedtuple( 58 'FilterShape2D', 'height, width, in_channels, out_channels') 59FilterShape2DTranspose = collections.namedtuple( 60 'FilterShape2DTranspose', 'height, width, out_channels, in_channels') 61 62LayerShapeNCDHW = collections.namedtuple( 63 'LayerShapeNCDHW', 'batch, channels, depth, height, width') 64FilterShape3D = collections.namedtuple( 65 'FilterShape3D', 'depth, height, width, in_channels, out_channels') 66 67 68class ConvolutionTest(test.TestCase): 69 """Tests for deterministic cuDNN functionality.""" 70 71 def _random_data_op(self, shape): 72 # np.random.random_sample can properly interpret either tf.TensorShape or 73 # namedtuple as a list. 74 return constant_op.constant( 75 2 * np.random.random_sample(shape) - 1, dtype=dtypes.float32) 76 77 def _random_out_op(self, in_shape, filter_shape, strides, padding, dilations): 78 # Choosing not to use array_op.zeros() to prevent possible removal by 79 # optimization 80 in_op = self._random_data_op(in_shape) 81 filter_op = self._random_data_op(filter_shape) 82 # Use the forward op's shape-inference 83 conv_op = nn_ops.conv2d( 84 in_op, filter_op, strides=strides, padding=padding, dilations=dilations) 85 out_shape = conv_op.get_shape() 86 out_op = self._random_data_op(out_shape) 87 return out_op 88 89 def _assert_reproducible(self, operation): 90 with test_util.force_gpu(): 91 result_1 = operation() 92 result_2 = operation() 93 self.assertAllEqual(result_1, result_2) 94 95 # The default forward algorithm choice, when using cuDNN 7, does not support 96 # the following layer configuration. This test case intends to confirm that 97 # an alternative algorithm is selected. Note that, in cuDNN 7, all forward 98 # algorithms are determnistic. 99 @test_util.run_cuda_only 100 def testConvForwardDefaultAlgorithmChoice(self): 101 in_shape = LayerShapeNCDHW(batch=2, channels=3, depth=5, height=7, width=6) 102 filter_shape = FilterShape3D( 103 depth=3, height=3, width=3, in_channels=3, out_channels=2) 104 in_op = self._random_data_op(in_shape) 105 filter_op = self._random_data_op(filter_shape) 106 self._assert_reproducible(lambda: nn_ops.conv3d( 107 in_op, 108 filter_op, 109 strides=[1, 1, 1, 1, 1], 110 padding='VALID', 111 data_format='NCDHW', 112 dilations=[1, 1, 2, 2, 2])) 113 114 # This test is primarily testing XLA since cuDNN forward convolutions are 115 # always deterministic, even when determinism is not enabled. The convolution 116 # configuration tested is nondeterministic with XLA when determinism is not 117 # enabled. 118 @test_util.run_cuda_only 119 def testConvForwardXLA(self): 120 in_shape = LayerShapeNCDHW( 121 batch=2, channels=8, depth=5, height=12, width=15) 122 filter_shape = FilterShape3D( 123 depth=3, height=3, width=3, in_channels=8, out_channels=1) 124 in_op = self._random_data_op(in_shape) 125 filter_op = self._random_data_op(filter_shape) 126 self._assert_reproducible(lambda: nn_ops.conv3d( 127 in_op, 128 filter_op, 129 strides=[1, 1, 1, 1, 1], 130 padding='VALID', 131 data_format='NCDHW', 132 dilations=[1, 1, 2, 2, 2])) 133 134 @test_util.run_cuda_only 135 def testConvBackwardFilterGradient(self, rate=1): 136 in_shape = LayerShapeNHWC(batch=8, height=64, width=64, channels=8) 137 filter_shape = FilterShape2D( 138 height=3, width=3, in_channels=8, out_channels=8) 139 in_op = self._random_data_op(in_shape) 140 strides = [1, 1, 1, 1] 141 padding = 'SAME' 142 dilations = [1, rate, rate, 1] 143 out_op = self._random_out_op(in_shape, filter_shape, strides, padding, 144 dilations) 145 self._assert_reproducible(lambda: nn_ops.conv2d_backprop_filter( 146 in_op, 147 filter_shape, 148 out_op, 149 strides=strides, 150 padding=padding, 151 dilations=dilations)) 152 153 # A configuration for this test could not be found that exercises 154 # nondeterminism when using XLA with determinism not enabled. 155 @test_util.run_cuda_only 156 def testConvBackwardFilterGradientWithDilations(self): 157 self.testConvBackwardFilterGradient(rate=2) 158 159 @test_util.run_cuda_only 160 def testConvBackwardInputGradient(self, rate=1): 161 in_shape = LayerShapeNHWC(batch=1, height=16, width=16, channels=1) 162 filter_shape = FilterShape2D( 163 height=7, width=7, in_channels=1, out_channels=3) 164 filter_op = self._random_data_op(filter_shape) 165 strides = [1, 1, 1, 1] 166 padding = 'SAME' 167 dilations = [1, rate, rate, 1] 168 out_op = self._random_out_op(in_shape, filter_shape, strides, padding, 169 dilations) 170 self._assert_reproducible(lambda: nn_ops.conv2d_backprop_input( 171 in_shape, 172 filter_op, 173 out_op, 174 strides=strides, 175 padding=padding, 176 dilations=dilations)) 177 178 # A configuration for this test could not be found that exercises 179 # nondeterminism when using XLA with determinism not enabled. 180 @test_util.run_cuda_only 181 def testConvBackwardInputGradientWithDilations(self): 182 self.testConvBackwardInputGradient(rate=2) 183 184 @test_util.run_cuda_only 185 def testConvTransposeForward(self, rate=1): 186 in_channels = 3 187 out_channels = 1 188 in_shape = LayerShapeNHWC( 189 batch=1, height=16, width=16, channels=in_channels) 190 filter_shape = FilterShape2DTranspose( 191 height=7, width=7, out_channels=out_channels, in_channels=in_channels) 192 in_op = self._random_data_op(in_shape) 193 filter_op = self._random_data_op(filter_shape) 194 out_shape = LayerShapeNHWC( 195 batch=in_shape.batch, 196 height=in_shape.height, 197 width=in_shape.width, 198 channels=out_channels) 199 self._assert_reproducible(lambda: nn_ops.conv2d_transpose_v2( 200 in_op, 201 filter_op, 202 out_shape, 203 strides=1, 204 padding='SAME', 205 data_format='NHWC', 206 dilations=[1, rate, rate, 1])) 207 208 # A configuration for this test could not be found that exercises 209 # nondeterminism when using XLA with determinism not enabled. 210 @test_util.run_cuda_only 211 def testConvTransposeForwardWithDilations(self): 212 self.testConvTransposeForward(rate=2) 213 214 @test_util.run_cuda_only 215 def testConvTransposeBackwardFilterGradient(self, rate=1): 216 in_channels = 8 217 out_channels = 8 218 in_shape = LayerShapeNHWC( 219 batch=8, height=64, width=64, channels=in_channels) 220 filter_shape = FilterShape2DTranspose( 221 height=3, width=3, out_channels=out_channels, in_channels=in_channels) 222 in_op = self._random_data_op(in_shape) 223 filter_op = self._random_data_op(filter_shape) 224 out_shape = LayerShapeNHWC( 225 batch=in_shape.batch, 226 height=in_shape.height, 227 width=in_shape.width, 228 channels=out_channels) 229 upstream_gradients = self._random_data_op(out_shape) 230 231 def gradient(): 232 with backprop.GradientTape() as tape: 233 tape.watch(filter_op) 234 op_output = nn_ops.conv2d_transpose_v2( 235 in_op, 236 filter_op, 237 out_shape, 238 strides=1, 239 padding='SAME', 240 data_format='NHWC', 241 dilations=[1, rate, rate, 1]) 242 gradient_injector_output = op_output * upstream_gradients 243 return tape.gradient(gradient_injector_output, [filter_op])[0] 244 245 self._assert_reproducible(gradient) 246 247 # A configuration for this test could not be found that exercises 248 # nondeterminism when using XLA with determinism not enabled. 249 @test_util.run_cuda_only 250 def testConvTransposeBackwardFilterGradientWithDilations(self): 251 self.testConvTransposeBackwardFilterGradient(rate=2) 252 253 # A configuration for this test could not be found that exercises 254 # nondeterminism when determinism is not enabled (for either XLA or non-XLA). 255 @test_util.run_cuda_only 256 def testConvTransposeBackwardInputGradient(self, rate=1): 257 in_channels = 1 258 out_channels = 3 259 in_shape = LayerShapeNHWC( 260 batch=1, height=16, width=16, channels=in_channels) 261 filter_shape = FilterShape2DTranspose( 262 height=7, width=7, out_channels=out_channels, in_channels=in_channels) 263 in_op = self._random_data_op(in_shape) 264 filter_op = self._random_data_op(filter_shape) 265 out_shape = LayerShapeNHWC( 266 batch=in_shape.batch, 267 height=in_shape.height, 268 width=in_shape.width, 269 channels=out_channels) 270 upstream_gradients = self._random_data_op(out_shape) 271 272 def gradient(): 273 with backprop.GradientTape() as tape: 274 tape.watch(in_op) 275 op_output = nn_ops.conv2d_transpose_v2( 276 in_op, 277 filter_op, 278 out_shape, 279 strides=1, 280 padding='SAME', 281 data_format='NHWC', 282 dilations=[1, rate, rate, 1]) 283 gradient_injector_output = op_output * upstream_gradients 284 return tape.gradient(gradient_injector_output, [in_op])[0] 285 286 self._assert_reproducible(gradient) 287 288 # A configuration for this test could not be found that exercises 289 # nondeterminism when determinism is not enabled (for either XLA or non-XLA). 290 @test_util.run_cuda_only 291 def testConvTransposeBackwardInputGradientWithDilations(self): 292 self.testConvTransposeBackwardInputGradient(rate=2) 293