1# Copyright 2017 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 pooling operations.""" 16 17import numpy as np 18 19from tensorflow.compiler.tests import xla_test 20from tensorflow.python.framework import dtypes 21from tensorflow.python.framework import ops 22from tensorflow.python.ops import array_ops 23from tensorflow.python.ops import gen_nn_ops 24from tensorflow.python.ops import nn_ops 25from tensorflow.python.platform import googletest 26 27 28def NHWCToNCHW(input_tensor): 29 """Convert the input from NHWC format to NCHW. 30 31 Args: 32 input_tensor: a 4-D tensor, or a 4-element array representing the same. 33 34 Returns: 35 the converted tensor or a shape array 36 """ 37 if isinstance(input_tensor, ops.Tensor): 38 return array_ops.transpose(input_tensor, [0, 3, 1, 2]) 39 else: 40 return [input_tensor[0], input_tensor[3], input_tensor[1], input_tensor[2]] 41 42 43def NCHWToNHWC(input_tensor): 44 """Convert the input from NCHW format to NHWC. 45 46 Args: 47 input_tensor: a 4-D tensor, or a 4-element array representing the same. 48 49 Returns: 50 the converted tensor or a shape array 51 """ 52 if isinstance(input_tensor, ops.Tensor): 53 return array_ops.transpose(input_tensor, [0, 2, 3, 1]) 54 else: 55 return [input_tensor[0], input_tensor[2], input_tensor[3], input_tensor[1]] 56 57 58def GetTestConfigs(): 59 """Get all the valid tests configs to run. 60 61 Returns: 62 all the valid test configs 63 """ 64 test_configs = ["NHWC", "NCHW"] 65 return test_configs 66 67 68class PoolingTest(xla_test.XLATestCase): 69 70 def _VerifyOneTest(self, pool_func, input_sizes, ksize, strides, padding, 71 data_format, expected): 72 """Verifies the output values of the pooling function. 73 74 Args: 75 pool_func: Function to be called, currently only co.MaxPool. 76 input_sizes: Input tensor dimensions. 77 ksize: The kernel size dimensions 78 strides: The stride dimensions 79 padding: Padding type. 80 data_format: The data format we use to run the pooling operation. 81 expected: An array containing the expected operation outputs. 82 """ 83 total_size = np.prod(input_sizes) 84 # Initializes the input tensor with array containing incrementing 85 # numbers from 1. 86 x = np.array([f * 1.0 for f in range(1, total_size + 1)], dtype=np.float32) 87 x = x.reshape(input_sizes) 88 with self.session() as sess: 89 with self.test_scope(): 90 inputs = array_ops.placeholder(dtypes.float32) 91 t = inputs 92 if data_format == "NCHW": 93 t = NHWCToNCHW(t) 94 ksize = NHWCToNCHW(ksize) 95 strides = NHWCToNCHW(strides) 96 t = pool_func(t, 97 ksize=ksize, 98 strides=strides, 99 padding=padding, 100 data_format=data_format) 101 if data_format == "NCHW": 102 t = NCHWToNHWC(t) 103 actual = sess.run(t, {inputs: x}) 104 self.assertAllClose(expected, actual.flatten(), rtol=1e-5, atol=1e-6) 105 106 def _VerifyValues(self, pool_func, input_sizes, ksize, strides, padding, 107 expected): 108 """Verifies the output values of the pooling function. 109 110 Args: 111 pool_func: Function to be called, co.MaxPool, co.AvgPool, 112 or the Lua version. 113 input_sizes: Input tensor dimensions. 114 ksize: The kernel size dimensions 115 strides: The stride dimensions 116 padding: Padding type. 117 expected: An array containing the expected operation outputs. 118 """ 119 for data_format in GetTestConfigs(): 120 self._VerifyOneTest(pool_func, input_sizes, ksize, strides, padding, 121 data_format, expected) 122 123 def testMaxPoolValidPadding(self): 124 expected_output = [13.0, 14.0, 15.0] 125 self._VerifyValues(nn_ops.max_pool, 126 input_sizes=[1, 3, 3, 3], 127 ksize=[1, 2, 2, 1], 128 strides=[1, 2, 2, 1], 129 padding="VALID", 130 expected=expected_output) 131 132 def testMaxPoolSamePadding(self): 133 expected_output = [13.0, 14.0, 15.0, 16.0, 17.0, 18.0] 134 self._VerifyValues(nn_ops.max_pool, 135 input_sizes=[1, 2, 3, 3], 136 ksize=[1, 2, 2, 1], 137 strides=[1, 2, 2, 1], 138 padding="SAME", 139 expected=expected_output) 140 141 def testMaxPoolSamePaddingNonSquareWindow(self): 142 # input is: 143 # [1.0, 2.0 144 # 3.0 4.0] 145 # 146 # Window of [x, x] should do: 147 # 148 # [max(1.0, 2.0), max(2.0, padded0), 149 # max(3.0, 4.0), max(4.0, padded0)] 150 self._VerifyValues( 151 nn_ops.max_pool, 152 input_sizes=[1, 2, 2, 1], 153 ksize=[1, 1, 2, 1], 154 strides=[1, 1, 1, 1], 155 padding="SAME", 156 expected=[2.0, 2.0, 4.0, 4.0]) 157 158 def testMaxPoolValidPaddingUnevenStride(self): 159 self._VerifyValues( 160 nn_ops.max_pool, 161 input_sizes=[1, 4, 4, 1], 162 ksize=[1, 2, 2, 1], 163 strides=[1, 1, 2, 1], 164 padding="VALID", 165 expected=[6.0, 8.0, 10.0, 12.0, 14.0, 16.0]) 166 self._VerifyValues( 167 nn_ops.max_pool, 168 input_sizes=[1, 4, 4, 1], 169 ksize=[1, 2, 2, 1], 170 strides=[1, 2, 1, 1], 171 padding="VALID", 172 expected=[6.0, 7.0, 8.0, 14.0, 15.0, 16.0]) 173 174 def testMaxPoolSamePaddingFilter4(self): 175 expected_output = [ 176 21.0, 22.0, 23.0, 24.0, 29.0, 30.0, 31.0, 32.0, 53.0, 54.0, 55.0, 56.0, 177 61.0, 62.0, 63.0, 64.0 178 ] 179 self._VerifyValues( 180 nn_ops.max_pool, 181 input_sizes=[1, 4, 4, 4], 182 ksize=[1, 2, 2, 1], 183 strides=[1, 2, 2, 1], 184 padding="SAME", 185 expected=expected_output) 186 187 def testMaxPoolSamePaddingFilter8(self): 188 expected_output = [ 189 145.0, 146.0, 147.0, 148.0, 149.0, 150.0, 151.0, 152.0, 161.0, 162.0, 190 163.0, 164.0, 165.0, 166.0, 167.0, 168.0, 177.0, 178.0, 179.0, 180.0, 191 181.0, 182.0, 183.0, 184.0, 185.0, 186.0, 187.0, 188.0, 189.0, 190.0, 192 191.0, 192.0, 273.0, 274.0, 275.0, 276.0, 277.0, 278.0, 279.0, 280.0, 193 289.0, 290.0, 291.0, 292.0, 293.0, 294.0, 295.0, 296.0, 305.0, 306.0, 194 307.0, 308.0, 309.0, 310.0, 311.0, 312.0, 313.0, 314.0, 315.0, 316.0, 195 317.0, 318.0, 319.0, 320.0, 401.0, 402.0, 403.0, 404.0, 405.0, 406.0, 196 407.0, 408.0, 417.0, 418.0, 419.0, 420.0, 421.0, 422.0, 423.0, 424.0, 197 433.0, 434.0, 435.0, 436.0, 437.0, 438.0, 439.0, 440.0, 441.0, 442.0, 198 443.0, 444.0, 445.0, 446.0, 447.0, 448.0, 465.0, 466.0, 467.0, 468.0, 199 469.0, 470.0, 471.0, 472.0, 481.0, 482.0, 483.0, 484.0, 485.0, 486.0, 200 487.0, 488.0, 497.0, 498.0, 499.0, 500.0, 501.0, 502.0, 503.0, 504.0, 201 505.0, 506.0, 507.0, 508.0, 509.0, 510.0, 511.0, 512.0 202 ] 203 self._VerifyValues( 204 nn_ops.max_pool, 205 input_sizes=[1, 8, 8, 8], 206 ksize=[1, 3, 3, 1], 207 strides=[1, 2, 2, 1], 208 padding="SAME", 209 expected=expected_output) 210 211 # Tests for DepthwiseMaxPooling on CPU only. 212 def testDepthwiseMaxPool1x1DepthWindow1(self): 213 # input is: 214 # [1.0, ..., 10.0] along depth, 215 # 216 # We maxpool by depth in patches of 2. 217 self._VerifyValues( 218 nn_ops.max_pool, 219 input_sizes=[1, 1, 1, 10], 220 ksize=[1, 1, 1, 2], 221 strides=[1, 1, 1, 2], 222 padding="SAME", 223 expected=[2.0, 4.0, 6.0, 8.0, 10.0]) 224 225 def testDepthwiseMaxPool2x2DepthWindow3(self): 226 # input is: 227 # 228 # a 2x2x6 cube, and we depthwise max across 3 to produce a 2x2x2 229 # output. Each node has contiguous values, so the depthwise max 230 # should be multiples of 3.0. 231 self._VerifyValues( 232 nn_ops.max_pool, 233 input_sizes=[1, 2, 2, 6], 234 ksize=[1, 1, 1, 3], 235 strides=[1, 1, 1, 3], 236 padding="SAME", 237 expected=[3.0, 6.0, 9.0, 12.0, 15.0, 18.0, 21.0, 24.0]) 238 239 def testKernelSmallerThanStrideValid(self): 240 self._VerifyValues( 241 nn_ops.max_pool, 242 input_sizes=[1, 7, 7, 1], 243 ksize=[1, 2, 2, 1], 244 strides=[1, 3, 3, 1], 245 padding="VALID", 246 expected=[9, 12, 30, 33]) 247 248 def testKernelSmallerThanStrideSame(self): 249 self._VerifyValues( 250 nn_ops.max_pool, 251 input_sizes=[1, 3, 3, 1], 252 ksize=[1, 1, 1, 1], 253 strides=[1, 2, 2, 1], 254 padding="SAME", 255 expected=[1, 3, 7, 9]) 256 257 self._VerifyValues( 258 nn_ops.max_pool, 259 input_sizes=[1, 4, 4, 1], 260 ksize=[1, 1, 1, 1], 261 strides=[1, 2, 2, 1], 262 padding="SAME", 263 expected=[1, 3, 9, 11]) 264 265 # Average pooling 266 def testAvgPoolValidPadding(self): 267 expected_output = [7, 8, 9] 268 self._VerifyValues( 269 nn_ops.avg_pool, 270 input_sizes=[1, 3, 3, 3], 271 ksize=[1, 2, 2, 1], 272 strides=[1, 2, 2, 1], 273 padding="VALID", 274 expected=expected_output) 275 276 def testAvgPoolSamePadding(self): 277 expected_output = [7., 8., 9., 11.5, 12.5, 13.5] 278 self._VerifyValues( 279 nn_ops.avg_pool, 280 input_sizes=[1, 2, 3, 3], 281 ksize=[1, 2, 2, 1], 282 strides=[1, 2, 2, 1], 283 padding="SAME", 284 expected=expected_output) 285 286 287class PoolGradTest(xla_test.XLATestCase): 288 289 CPU_DEVICE = "/job:localhost/replica:0/task:0/cpu:0" 290 291 def _VerifyOneTest(self, 292 pool_func, 293 pool_grad_func, 294 input_sizes, 295 ksize, 296 strides, 297 padding, 298 data_format, 299 pool_grad_grad_func=None): 300 """Verifies the output values of the pooling gradient function. 301 302 Args: 303 pool_func: Forward pooling function 304 pool_grad_func: Pooling gradient function for pool_grad_func 305 input_sizes: Input tensor dimensions. 306 ksize: The kernel size dimensions 307 strides: The stride dimensions 308 padding: Padding type. 309 data_format: The data format we use to run the pooling operation. 310 pool_grad_grad_func: Second-order gradient function, if available. 311 """ 312 total_size = np.prod(input_sizes) 313 # TODO(b/73062247): MaxPoolGradGrad can confuse gradients when x is equally 314 # maximal at 16 bits. Switch to np.random.randn when resolved. 315 x = np.arange(1, total_size + 1, dtype=np.float32) 316 x *= (np.random.randint(2, size=total_size) * 2 - 1) # Flip signs randomly 317 # Verify some specifically interesting values... 318 x[np.random.choice(total_size)] = np.inf 319 x[np.random.choice(total_size)] = -np.inf 320 # TODO(b/74222344): Fix nan handling for max pool grad. 321 # x[np.random.choice(total_size)] = np.nan 322 x = x.reshape(input_sizes) 323 with self.session() as sess: 324 # Use the forward pool function to compute some corresponding outputs 325 # (needed for the CPU device, and we need the shape in both cases). 326 with ops.device(self.CPU_DEVICE): 327 inputs = array_ops.placeholder(dtypes.float32, shape=input_sizes) 328 outputs = pool_func( 329 inputs, 330 ksize=ksize, 331 strides=strides, 332 padding=padding, 333 data_format="NHWC") 334 335 output_vals = np.array(sess.run(outputs, {inputs: x})) 336 output_gradient_vals = np.arange( 337 1, output_vals.size + 1, dtype=np.float32) 338 output_gradient_vals = output_gradient_vals.reshape(output_vals.shape) 339 output_grad_grad_vals = np.arange(1, x.size + 1, dtype=np.float32) 340 output_grad_grad_vals = output_grad_grad_vals.reshape(x.shape) 341 342 # Use the Tensorflow CPU pooling gradient to compute the expected input 343 # gradients. 344 with ops.device(self.CPU_DEVICE): 345 output_gradients = array_ops.placeholder( 346 dtypes.float32, shape=output_vals.shape) 347 expected_input_gradients = pool_grad_func( 348 inputs, 349 outputs, 350 output_gradients, 351 ksize=ksize, 352 strides=strides, 353 padding=padding, 354 data_format="NHWC") 355 expected_input_gradient_vals = sess.run( 356 expected_input_gradients, 357 {inputs: x, 358 output_gradients: output_gradient_vals}) 359 360 output_grad_gradients = array_ops.placeholder( 361 dtypes.float32, shape=expected_input_gradient_vals.shape) 362 if pool_grad_grad_func is not None: 363 expected_grad_gradients = pool_grad_grad_func( 364 inputs, 365 outputs, 366 output_grad_gradients, 367 ksize=ksize, 368 strides=strides, 369 padding=padding, 370 data_format="NHWC") 371 expected_grad_gradients_vals = sess.run(expected_grad_gradients, { 372 inputs: x, 373 output_grad_gradients: output_grad_grad_vals 374 }) 375 376 # Run the gradient op on the XLA device 377 with self.test_scope(): 378 outputs = array_ops.placeholder(dtypes.float32, shape=output_vals.shape) 379 xla_inputs = inputs 380 xla_outputs = outputs 381 xla_output_gradients = output_gradients 382 xla_output_grad_gradients = output_grad_gradients 383 xla_ksize = ksize 384 xla_strides = strides 385 if data_format == "NCHW": 386 xla_inputs = NHWCToNCHW(inputs) 387 xla_outputs = NHWCToNCHW(outputs) 388 xla_output_gradients = NHWCToNCHW(output_gradients) 389 xla_output_grad_gradients = NHWCToNCHW(output_grad_gradients) 390 xla_ksize = NHWCToNCHW(ksize) 391 xla_strides = NHWCToNCHW(strides) 392 actual_input_gradients = pool_grad_func( 393 xla_inputs, 394 xla_outputs, 395 xla_output_gradients, 396 ksize=xla_ksize, 397 strides=xla_strides, 398 padding=padding, 399 data_format=data_format) 400 if data_format == "NCHW": 401 actual_input_gradients = NCHWToNHWC(actual_input_gradients) 402 if pool_grad_grad_func is not None: 403 actual_grad_gradients = pool_grad_grad_func( 404 xla_inputs, 405 xla_outputs, 406 xla_output_grad_gradients, 407 ksize=xla_ksize, 408 strides=xla_strides, 409 padding=padding, 410 data_format=data_format) 411 if data_format == "NCHW": 412 actual_grad_gradients = NCHWToNHWC(actual_grad_gradients) 413 actual_input_gradients_vals = sess.run(actual_input_gradients, { 414 inputs: x, 415 outputs: output_vals, 416 output_gradients: output_gradient_vals 417 }) 418 # Compare the Tensorflow and XLA results. 419 self.assertAllClose( 420 expected_input_gradient_vals, 421 actual_input_gradients_vals, 422 rtol=1e-4, 423 atol=1e-6) 424 self.assertShapeEqual(actual_input_gradients_vals, inputs) 425 426 if pool_grad_grad_func is not None: 427 actual_grad_gradients_vals = sess.run( 428 actual_grad_gradients, { 429 inputs: x, 430 outputs: output_vals, 431 output_grad_gradients: output_grad_grad_vals 432 }) 433 434 # Compare the Tensorflow and XLA results. 435 self.assertAllClose( 436 expected_grad_gradients_vals, 437 actual_grad_gradients_vals, 438 rtol=1e-4, 439 atol=1e-6) 440 self.assertShapeEqual(actual_grad_gradients_vals, outputs) 441 442 def _VerifyValues(self, 443 pool_func, 444 pool_grad_func, 445 input_sizes, 446 ksize, 447 strides, 448 padding, 449 pool_grad_grad_func=None): 450 """Verifies the output values of the pooling function. 451 452 Args: 453 pool_func: Pooling function to be called, e.g., tf.nn.max_pool2d 454 pool_grad_func: Corresponding pooling gradient function. 455 input_sizes: Input tensor dimensions. 456 ksize: The kernel size dimensions 457 strides: The stride dimensions 458 padding: Padding type. 459 pool_grad_grad_func: Second-order gradient function, if available. 460 """ 461 for data_format in GetTestConfigs(): 462 self._VerifyOneTest( 463 pool_func, 464 pool_grad_func, 465 input_sizes, 466 ksize, 467 strides, 468 padding, 469 data_format, 470 pool_grad_grad_func=pool_grad_grad_func) 471 472 def _TestPooling(self, forward_op, backward_op, pool_grad_grad_func=None): 473 # VALID padding 474 self._VerifyValues( 475 forward_op, 476 backward_op, 477 input_sizes=[1, 3, 3, 3], 478 ksize=[1, 2, 2, 1], 479 strides=[1, 2, 2, 1], 480 padding="VALID", 481 pool_grad_grad_func=pool_grad_grad_func) 482 483 # SAME padding 484 self._VerifyValues( 485 forward_op, 486 backward_op, 487 input_sizes=[1, 2, 3, 3], 488 ksize=[1, 2, 2, 1], 489 strides=[1, 2, 2, 1], 490 padding="SAME", 491 pool_grad_grad_func=pool_grad_grad_func) 492 493 # SAME padding, non square window 494 self._VerifyValues( 495 forward_op, 496 backward_op, 497 input_sizes=[1, 2, 2, 1], 498 ksize=[1, 1, 2, 1], 499 strides=[1, 1, 1, 1], 500 padding="SAME", 501 pool_grad_grad_func=pool_grad_grad_func) 502 503 # VALID padding, uneven stride 504 self._VerifyValues( 505 forward_op, 506 backward_op, 507 input_sizes=[1, 4, 4, 1], 508 ksize=[1, 2, 2, 1], 509 strides=[1, 1, 2, 1], 510 padding="VALID", 511 pool_grad_grad_func=pool_grad_grad_func) 512 self._VerifyValues( 513 forward_op, 514 backward_op, 515 input_sizes=[1, 4, 4, 1], 516 ksize=[1, 2, 2, 1], 517 strides=[1, 2, 1, 1], 518 padding="VALID", 519 pool_grad_grad_func=pool_grad_grad_func) 520 521 # SAME padding, size 4 input 522 self._VerifyValues( 523 forward_op, 524 backward_op, 525 input_sizes=[1, 4, 4, 4], 526 ksize=[1, 2, 2, 1], 527 strides=[1, 2, 2, 1], 528 padding="SAME", 529 pool_grad_grad_func=pool_grad_grad_func) 530 531 # SAME padding, size 8 input 532 self._VerifyValues( 533 forward_op, 534 backward_op, 535 input_sizes=[1, 8, 8, 8], 536 ksize=[1, 3, 3, 1], 537 strides=[1, 2, 2, 1], 538 padding="SAME", 539 pool_grad_grad_func=pool_grad_grad_func) 540 541 def testMaxPool(self): 542 self._TestPooling( 543 nn_ops.max_pool, 544 gen_nn_ops.max_pool_grad, 545 pool_grad_grad_func=gen_nn_ops.max_pool_grad_grad) 546 547 def testAvgPool(self): 548 # Wrapper around AvgPoolGrad that ignores extra arguments needed by 549 # MaxPoolGrad. 550 def AvgPoolGrad(inputs, outputs, output_gradients, ksize, strides, padding, 551 data_format): 552 del outputs # Unused by average-pooling gradients. 553 return gen_nn_ops.avg_pool_grad( 554 inputs.get_shape().as_list(), 555 output_gradients, 556 ksize=ksize, 557 strides=strides, 558 padding=padding, 559 data_format=data_format) 560 561 self._TestPooling(nn_ops.avg_pool, AvgPoolGrad) 562 563 # The CPU implementation of AvgPoolGrad doesn't accept kernels smaller than 564 # the stride size, so we only run the following tests on MaxPoolGrad. 565 566 def testMaxPoolKernelSmallerThanStrideValid(self): 567 self._VerifyValues( 568 nn_ops.max_pool, 569 gen_nn_ops.max_pool_grad, 570 input_sizes=[1, 7, 7, 1], 571 ksize=[1, 2, 2, 1], 572 strides=[1, 3, 3, 1], 573 padding="VALID") 574 575 def testMaxPoolKernelSmallerThanStrideSame(self): 576 self._VerifyValues( 577 nn_ops.max_pool, 578 gen_nn_ops.max_pool_grad, 579 input_sizes=[1, 3, 3, 1], 580 ksize=[1, 1, 1, 1], 581 strides=[1, 2, 2, 1], 582 padding="SAME") 583 584 self._VerifyValues( 585 nn_ops.max_pool, 586 gen_nn_ops.max_pool_grad, 587 input_sizes=[1, 4, 4, 1], 588 ksize=[1, 1, 1, 1], 589 strides=[1, 2, 2, 1], 590 padding="SAME") 591 592 593if __name__ == "__main__": 594 googletest.main() 595