Where actually is max_pool_with_argmax implemented?

jondo · March 15, 2022, 2:50am

Hello,

in current Tensorflow 2.8.0 nn.max_pool_with_argmax is defined in nn_ops.py as max_pool_with_argmax_v2, saying

The indices returned are always in [0, height) x [0, width) before flattening, even if padding is involved and the mathematically correct answer is outside (either negative or too large). This is a bug, but fixing it is difficult to do in a safe backwards compatible way, especially due to flattening.

In order to understand what is actually going on, I wanted to take a look at how the ‘eager’ version of this operation is implemented. For this, I saw that in gen_nn_ops.py's max_pool_with_argmax TFE_Py_FastPathExecute is called with an argument"MaxPoolWithArgmax". Now I am stuck: Where is this function call dispatched to? Where is the code that actually executes the max pooling?

Bhack · March 15, 2022, 2:54pm

I think these are the CPU and GPU c++/cuda kernels:

github.com

tensorflow/tensorflow/blob/v2.8.0/tensorflow/core/kernels/maxpooling_op.cc

/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

// See docs in ../ops/nn_ops.cc.

#define EIGEN_USE_THREADS

#include "tensorflow/core/kernels/maxpooling_op.h"

This file has been truncated. show original

github.com

tensorflow/tensorflow/blob/v2.8.0/tensorflow/core/kernels/maxpooling_op_gpu.cu.cc

/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM

#define EIGEN_USE_GPU

#include <stdio.h>

This file has been truncated. show original