Function MaxPool3dGrad

picasso/point/pi_modules/source/pcloud_pool3d.cpp:57–83

torch::Tensor MaxPool3dGrad(
    torch::Tensor grad_output,   // gradient of pooled features: concat_Mp * in_channels
    torch::Tensor input,         // input features: concat_Np * in_channels
    torch::Tensor max_index)     // the neighbor gives maximum activation: concat_Mp * in_channels
{
    CHECK_INPUT(input,2);
    CHECK_INPUT(grad_output,2);
    CHECK_INPUT(max_index,2);

    // get the dims required by computations
    int Np = input.size(0);       // number of input points
    int C  = input.size(1);       // number of input channels
    int Mp = grad_output.size(0); // number of output points

    TORCH_CHECK(max_index.dim()==2, "rank of max_index should be 2, i.e. (Mp, in_channels)");

    // get the c++ pointers to the input tensors
    const float* gradOut_ptr = grad_output.data_ptr<float>();
    const int* maxIndex_ptr  = max_index.data_ptr<int32_t>();

    // create an output tensor
    auto grad_input = torch::zeros({Np,C}, input.options());
    float* gradIn_ptr = grad_input.data_ptr<float>();

    maxPool3dGradLauncher(Np, Mp, C, maxIndex_ptr, gradOut_ptr, gradIn_ptr);
    return grad_input;
}


torch::Tensor AvgPool3d(