Method Compute

tensorflow/core/kernels/avgpooling_op.cc:253–362

  }

  void Compute(OpKernelContext* context) override {
    const Tensor& tensor_in_shape = context->input(0);
    const Tensor& out_backprop = context->input(1);
    // For avgpooling, tensor_in_shape should have 1 dimension, and 4 elements.
    OP_REQUIRES(
        context,
        tensor_in_shape.dims() == 1 && tensor_in_shape.NumElements() == 4,
        errors::InvalidArgument("out_backprop must be 1-dimensional and 4 "
                                "elements"));
    // For avgpooling, out_backprop should have 4 dimensions.
    OP_REQUIRES(context, out_backprop.dims() == 4,
                errors::InvalidArgument("out_backprop must be 4-dimensional"));
    const int64 out_backprop_batch = out_backprop.dim_size(0);
    const int64 out_backprop_rows = out_backprop.dim_size(1);
    const int64 out_backprop_cols = out_backprop.dim_size(2);
    const int64 out_backprop_depth = out_backprop.dim_size(3);

    TensorShape output_shape;
    auto shape_vec = tensor_in_shape.vec<int32>();
    for (int64 i = 0; i < tensor_in_shape.NumElements(); ++i) {
      output_shape.AddDim(shape_vec(i));
    }
    const int64 in_rows = output_shape.dim_size(1);
    const int64 in_cols = output_shape.dim_size(2);

    Tensor* output = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output));
    output->flat<T>().setZero();

    const int window_rows = ksize_[1];
    const int window_cols = ksize_[2];
    const int depth_window = ksize_[3];

    const int row_stride = stride_[1];
    const int col_stride = stride_[2];

    // We (will) use different code for spatial pooling and
    // non-spatial pooling.
    //
    // Spatial pooling is when depth_window = 1
    OP_REQUIRES(context, depth_window == 1,
                errors::Unimplemented("Non-spatial pooling is not "
                                      "yet supported. Volunteers? :)"));

    int64 out_height, out_width, pad_rows, pad_cols;
    OP_REQUIRES_OK(context,
                   GetWindowedOutputSize(in_rows, window_rows, row_stride,
                                         padding_, &out_height, &pad_rows));
    OP_REQUIRES_OK(context,
                   GetWindowedOutputSize(in_cols, window_cols, col_stride,
                                         padding_, &out_width, &pad_cols));

    const T* out_backprop_ptr = out_backprop.flat<T>().data();
    T* input_backprop_ptr = output->flat<T>().data();

    auto shard = [context, out_backprop_ptr, input_backprop_ptr,
                  out_backprop_rows, out_backprop_cols, out_backprop_depth,
                  in_rows, in_cols, window_rows, window_cols, row_stride,