Repository: Project-MONAI/MONAI

Issue:
Additional features for the BilateralFilter layer
**Is your feature request related to a problem? Please describe.**
For using the BilateralFilter layer (by @charliebudd ) more flexibly, I would really appreciate two extensions to the current filter implementation.

**Describe the solution you'd like**
1. It would be nice to implement the backpropagation of a loss into the filter parameters (spatial and range sigma) to be able to learn them directly.
2. In addition, independent spatial filter kernel sizes (sigma) for different spatial dimensions would be very interesting to have (Gaussian filter kernel can have a different size in different dimensions).


Discussion:
Just to add to this, the PHLFilter may be useful to you as it can have varying spatial sigmas. To construct a gaussian filter just provide the spatial coordinates of each pixel as the feature tensor. For a billateral filter provide the spatial and color coordinates. 
Hi @charliebudd,
I implemented a fully differentiable bilateral filter layer (with respect to its input as well as to its spatial/range sigmas) based on an analytical derivative of the filter algorithm (both C++/cpu and Cuda implementation, yet only tested on Ubuntu): https://github.com/faebstn96/trainable-bilateral-filter-source
You can find its detailed mathematical description together with some experiments in our arXiv preprint: https://arxiv.org/pdf/2201.10345.pdf
As the implementation is based on the forward pass of the current MONAI bilateral filter implementation, an integration into MONAI should be fairly straight forward. I also included a gradcheck.py script, which confirms that the analytical gradient matches the numerical derivative of the forward pass using Pytorch's torch.autograd.gradcheck function. A unit test could be created from that script. Two simple examples are also included in the repository. Of course, the gradient updates of the kernel widths can be deactivated by simply removing the filter parameters from the Pytorch graph and the filter can be used as a standard differentiable bilateral filter with hyperparameters. I would be happy to contribute to the MONAI library with my implementation.

One more comment regarding the current bilateral filter layer implementation in the MONAI repo: Please note that the current backward pass of the bilateral filter (gradient with respect to the filter input) in the MONAI library does not calculate the correct derivative of the forward pass. In the current implementation, the backward pass is calculated with the same function as the forward pass (https://github.com/Project-MONAI/MONAI/blob/dev/monai/networks/layers/filtering.py). This may hold for a standard convolutional layer, but is not correct for the bilateral filter as the input image is additionally used as input to the range kernel, which complicates the derivative (see our arXiv paper). A correct analytical derivative is implemented in my layer and validated by the gradcheck test.

Please let me know what you think and if the layer is worth integrating into MONAI. I am happy to answer any questions or help if you have problems running the code from my repo.
This looks great! And good spot with the current monai filter, It should indeed be saving the input image for guiding the bilateral filter for the backwards step. I think it makes sense to put together a PR to update to use your analytical derivative.
Great! Let me know if/how I can help you with the PR.

Maybe one other comment here: At the moment the bilateral filter layer does not support channel dimensions >1. Currently, this will run into an assert. When implementing this, one could also allow defining the number of output channels like in conventional convolutional layers. However, this might be an enhancement for later. I did not find the time yet to properly implement that. The batch dimension is, however, already handled correctly.
Yes I noted that, we would also need to handle an input with 2 spatial dimensions but I dont see any particular problems here. I'm happy to implement this. 
Yes, I guess you could handle two spatial dimensions the same way as it is done in the current MONAI bilateral filter implementation. Thank you so much, let me know when you need help!
Hi @charliebudd,
Did you already have progress on the implementation of the bilateral filter layer? Please let me know if there are problems or you need help!
I have not had the capacity to work on this as of yet. 
Sure, no worries - just let me know when you need help!
To add here - Our arXiv paper is published open-access at Medical Physics (https://aapm.onlinelibrary.wiley.com/doi/10.1002/mp.15718). In principle the bilateral filter layer should be eligible to be added to the Project-MONAI/research-contributions, right? Should we do that @charliebudd?
Hi @charliebudd, any updates on this?
Best
Fabi
Thanks @faebstn96, have you compared the backward pass with a naive PyTorch implementation, would you like to join one of our weekly dev meetings to discuss this? let me send you an email. 
Thanks for your invitation, I would be happy to join!
For development purpose I implemented a naive forward pass in PyTorch and tried to make the filter parameters trainable via autograd. As expected the performance was extremely poor in terms of speed and memory requirement (although the implementation was not really optimized). I did not directly compare the PyTorch implementation with the custom analytical gradient-based layer though. To be sure that the gradient is correct my GitHub repository contains a gradcheck.py script that compares the implemented analytical with a numerical gradient using the torch.autograd.gradcheck function.

I will respond via email as well, thanks!

Resolving patch:
diff --git a/monai/networks/layers/__init__.py b/monai/networks/layers/__init__.py
--- a/monai/networks/layers/__init__.py
+++ b/monai/networks/layers/__init__.py
@@ -14,7 +14,7 @@
 from .convutils import calculate_out_shape, gaussian_1d, polyval, same_padding, stride_minus_kernel_padding
 from .drop_path import DropPath
 from .factories import Act, Conv, Dropout, LayerFactory, Norm, Pad, Pool, split_args
-from .filtering import BilateralFilter, PHLFilter
+from .filtering import BilateralFilter, PHLFilter, TrainableBilateralFilter
 from .gmm import GaussianMixtureModel
 from .simplelayers import (
     LLTM,
diff --git a/monai/networks/layers/filtering.py b/monai/networks/layers/filtering.py
--- a/monai/networks/layers/filtering.py
+++ b/monai/networks/layers/filtering.py
@@ -17,7 +17,7 @@
 
 _C, _ = optional_import("monai._C")
 
-__all__ = ["BilateralFilter", "PHLFilter"]
+__all__ = ["BilateralFilter", "PHLFilter", "TrainableBilateralFilter"]
 
 
 class BilateralFilter(torch.autograd.Function):
@@ -101,3 +101,150 @@ def backward(ctx, grad_output):
         # scaled_features, = ctx.saved_variables
         # grad_input = _C.phl_filter(grad_output, scaled_features)
         # return grad_input
+
+
+class TrainableBilateralFilterFunction(torch.autograd.Function):
+    """
+    torch.autograd.Function for the TrainableBilateralFilter layer.
+
+    See:
+        F. Wagner, et al., Ultralow-parameter denoising: Trainable bilateral filter layers in
+        computed tomography, Medical Physics (2022), https://doi.org/10.1002/mp.15718
+
+    Args:
+        input: input tensor to be filtered.
+
+        sigma x: trainable standard deviation of the spatial filter kernel in x direction.
+
+        sigma y: trainable standard deviation of the spatial filter kernel in y direction.
+
+        sigma z: trainable standard deviation of the spatial filter kernel in z direction.
+
+        color sigma: trainable standard deviation of the intensity range kernel. This filter
+            parameter determines the degree of edge preservation.
+
+    Returns:
+        output (torch.Tensor): filtered tensor.
+    """
+
+    @staticmethod
+    def forward(ctx, input_img, sigma_x, sigma_y, sigma_z, color_sigma):
+        output_tensor, output_weights_tensor, do_dx_ki, do_dsig_r, do_dsig_x, do_dsig_y, do_dsig_z = _C.tbf_forward(
+            input_img, sigma_x, sigma_y, sigma_z, color_sigma
+        )
+
+        ctx.save_for_backward(
+            input_img,
+            sigma_x,
+            sigma_y,
+            sigma_z,
+            color_sigma,
+            output_tensor,
+            output_weights_tensor,
+            do_dx_ki,
+            do_dsig_r,
+            do_dsig_x,
+            do_dsig_y,
+            do_dsig_z,
+        )
+
+        return output_tensor
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input_img = ctx.saved_tensors[0]  # input image
+        sigma_x = ctx.saved_tensors[1]
+        sigma_y = ctx.saved_tensors[2]
+        sigma_z = ctx.saved_tensors[3]
+        color_sigma = ctx.saved_tensors[4]
+        output_tensor = ctx.saved_tensors[5]  # filtered image
+        output_weights_tensor = ctx.saved_tensors[6]  # weights
+        do_dx_ki = ctx.saved_tensors[7]  # derivative of output with respect to input, while k==i
+        do_dsig_r = ctx.saved_tensors[8]  # derivative of output with respect to range sigma
+        do_dsig_x = ctx.saved_tensors[9]  # derivative of output with respect to sigma x
+        do_dsig_y = ctx.saved_tensors[10]  # derivative of output with respect to sigma y
+        do_dsig_z = ctx.saved_tensors[11]  # derivative of output with respect to sigma z
+
+        # calculate gradient with respect to the sigmas
+        grad_color_sigma = torch.sum(grad_output * do_dsig_r)
+        grad_sig_x = torch.sum(grad_output * do_dsig_x)
+        grad_sig_y = torch.sum(grad_output * do_dsig_y)
+        grad_sig_z = torch.sum(grad_output * do_dsig_z)
+
+        grad_output_tensor = _C.tbf_backward(
+            grad_output,
+            input_img,
+            output_tensor,
+            output_weights_tensor,
+            do_dx_ki,
+            sigma_x,
+            sigma_y,
+            sigma_z,
+            color_sigma,
+        )
+
+        return grad_output_tensor, grad_sig_x, grad_sig_y, grad_sig_z, grad_color_sigma
+
+
+class TrainableBilateralFilter(torch.nn.Module):
+    """
+    Implementation of a trainable bilateral filter layer as proposed in the corresponding publication.
+    All filter parameters can be trained data-driven. The spatial filter kernels x, y, and z determine
+    image smoothing whereas the color parameter specifies the amount of edge preservation.
+    Can run on 1D, 2D, or 3D tensors (on top of Batch and Channel dimensions).
+
+    See:
+        F. Wagner, et al., Ultralow-parameter denoising: Trainable bilateral filter layers in
+        computed tomography, Medical Physics (2022), https://doi.org/10.1002/mp.15718
+
+    Args:
+        input: input tensor to be filtered.
+
+        sigma x: trainable standard deviation of the spatial filter kernel in x direction.
+
+        sigma y: trainable standard deviation of the spatial filter kernel in y direction.
+
+        sigma z: trainable standard deviation of the spatial filter kernel in z direction.
+
+        sigma color: trainable standard deviation of the intensity range kernel. This filter
+            parameter determines the degree of edge preservation.
+
+    Returns:
+        output (torch.Tensor): filtered tensor.
+    """
+
+    def __init__(self, sigma_x, sigma_y, sigma_z, sigma_color):
+        super().__init__()
+
+        # Register sigmas as trainable parameters.
+        self.sigma_x = torch.nn.Parameter(torch.tensor(sigma_x))
+        self.sigma_y = torch.nn.Parameter(torch.tensor(sigma_y))
+        self.sigma_z = torch.nn.Parameter(torch.tensor(sigma_z))
+        self.sigma_color = torch.nn.Parameter(torch.tensor(sigma_color))
+
+    def forward(self, input_tensor):
+        assert input_tensor.shape[1] == 1, (
+            "Currently channel dimensions >1 are not supported. "
+            "Please use multiple parallel filter layers if you want "
+            "to filter multiple channels."
+        )
+
+        len_input = len(input_tensor.shape)
+
+        # C++ extension so far only supports 5-dim inputs.
+        if len_input == 3:
+            input_tensor = input_tensor.unsqueeze(3).unsqueeze(4)
+        elif len_input == 4:
+            input_tensor = input_tensor.unsqueeze(4)
+
+        prediction = TrainableBilateralFilterFunction.apply(
+            input_tensor, self.sigma_x, self.sigma_y, self.sigma_z, self.sigma_color
+        )
+
+        # Make sure to return tensor of the same shape as the input.
+        if len_input == 3:
+            prediction = prediction.squeeze(4).squeeze(3)
+        elif len_input == 4:
+            prediction = prediction.squeeze(4)
+
+        return prediction


Test patch:
diff --git a/tests/test_trainable_bilateral.py b/tests/test_trainable_bilateral.py
new file mode 100644
--- /dev/null
+++ b/tests/test_trainable_bilateral.py
@@ -0,0 +1,478 @@
+# Copyright (c) MONAI Consortium
+# 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.
+
+from __future__ import annotations
+
+import unittest
+
+import numpy as np
+import torch
+from parameterized import parameterized
+from torch.autograd import gradcheck
+
+from monai.networks.layers.filtering import TrainableBilateralFilterFunction
+from tests.utils import skip_if_no_cpp_extension, skip_if_no_cuda
+
+TEST_CASES = [
+    [
+        # Case Description
+        "1 dimension, 1 channel, low spatial sigmas, low color sigma",
+        # (sigma_x, sigma_y, sigma_z, color_sigma)
+        (1.0, 1.0, 1.0, 0.2),
+        # Input
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [1, 0, 0, 0, 1]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0, 0, 1, 0, 0]
+            ],
+        ],
+        # Expected
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [0.999997, 0.000001, 0.000000, 0.000001, 0.999997]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0.000000, 0.000001, 0.999995, 0.000001, 0.000000]
+            ],
+        ],
+    ],
+    [
+        # Case Description
+        "1 dimension, 1 channel, low spatial sigmas, high color sigma",
+        # (sigma_x, sigma_y, sigma_z, color_sigma)
+        (1, 1, 1, 0.9),
+        # Input
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [1, 0, 0, 0, 1]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0, 0, 1, 0, 0]
+            ],
+        ],
+        # Expected
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [0.714200, 0.158126, 0.061890, 0.158126, 0.714200]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0.043465, 0.158126, 0.555452, 0.158126, 0.043465]
+            ],
+        ],
+    ],
+    [
+        # Case Description
+        "1 dimension, 1 channel, high spatial sigmas, low color sigma",
+        # (sigma_x, sigma_y, sigma_z, color_sigma)
+        (4, 4, 4, 0.2),
+        # Input
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [1, 0, 0, 0, 1]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0, 0, 1, 0, 0]
+            ],
+        ],
+        # Expected
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [0.999994, 0.000002, 0.000002, 0.000002, 0.999994]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0.000001, 0.000001, 0.999986, 0.000001, 0.000001]
+            ],
+        ],
+    ],
+    [
+        # Case Description
+        "1 dimension, 1 channel, high spatial sigmas, high color sigma",
+        # (sigma_x, sigma_y, sigma_z, color_sigma)
+        (4, 4, 4, 0.9),
+        # Input
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [1, 0, 0, 0, 1]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0, 0, 1, 0, 0]
+            ],
+        ],
+        # Expected
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [0.533282, 0.245915, 0.244711, 0.245915, 0.533282]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0.125052, 0.126608, 0.333592, 0.126608, 0.125052]
+            ],
+        ],
+    ],
+    [
+        # Case Description
+        "2 dimensions, 1 channel, high spatial sigmas, high color sigma",
+        # (sigma_x, sigma_y, sigma_z, color_sigma)
+        (4, 4, 4, 0.9),
+        # Input
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [[1, 0, 0, 0, 1], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [1, 0, 0, 0, 1]]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]
+            ],
+        ],
+        # Expected
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [
+                    [0.239789, 0.082990, 0.082630, 0.082990, 0.239789],
+                    [0.082990, 0.081934, 0.081579, 0.081934, 0.082990],
+                    [0.082630, 0.081579, 0.081225, 0.081579, 0.082630],
+                    [0.082990, 0.081934, 0.081579, 0.081934, 0.082990],
+                    [0.239789, 0.082990, 0.082630, 0.082990, 0.239789],
+                ]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [
+                    [0.024155, 0.024432, 0.024525, 0.024432, 0.024155],
+                    [0.024432, 0.024712, 0.024806, 0.024712, 0.024432],
+                    [0.024525, 0.024806, 0.080686, 0.024806, 0.024525],
+                    [0.024432, 0.024712, 0.024806, 0.024712, 0.024432],
+                    [0.024155, 0.024432, 0.024525, 0.024432, 0.024155],
+                ]
+            ],
+        ],
+    ],
+    [
+        # Case Description
+        "3 dimensions, 1 channel, high spatial sigmas, high color sigma",
+        # (sigma_x, sigma_y, sigma_z, color_sigma)
+        (4, 4, 4, 0.9),
+        # Input
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [
+                    # Frame 0
+                    [[1, 0, 0, 0, 1], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [1, 0, 0, 0, 1]],
+                    # Frame 1
+                    [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]],
+                    # Frame 2
+                    [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]],
+                    # Frame 3
+                    [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]],
+                    # Frame 4
+                    [[1, 0, 0, 0, 1], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [1, 0, 0, 0, 1]],
+                ]
+            ]
+        ],
+        # Expected
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [
+                    # Frame 0
+                    [
+                        [0.098142, 0.030317, 0.030191, 0.030316, 0.098142],
+                        [0.030316, 0.029947, 0.029822, 0.029947, 0.030316],
+                        [0.030191, 0.029822, 0.029698, 0.029822, 0.030191],
+                        [0.030316, 0.029947, 0.029822, 0.029947, 0.030316],
+                        [0.098142, 0.030317, 0.030191, 0.030317, 0.098142],
+                    ],
+                    # Frame 1
+                    [
+                        [0.030316, 0.029947, 0.029822, 0.029947, 0.030316],
+                        [0.029947, 0.029581, 0.029458, 0.029581, 0.029947],
+                        [0.029822, 0.029458, 0.029336, 0.029458, 0.029822],
+                        [0.029947, 0.029581, 0.029458, 0.029581, 0.029947],
+                        [0.030316, 0.029947, 0.029822, 0.029947, 0.030316],
+                    ],
+                    # Frame 2
+                    [
+                        [0.030191, 0.029822, 0.029698, 0.029822, 0.030191],
+                        [0.029822, 0.029458, 0.029336, 0.029458, 0.029822],
+                        [0.029698, 0.029336, 0.029214, 0.029336, 0.029698],
+                        [0.029822, 0.029458, 0.029336, 0.029458, 0.029822],
+                        [0.030191, 0.029822, 0.029698, 0.029822, 0.030191],
+                    ],
+                    # Frame 3
+                    [
+                        [0.030316, 0.029947, 0.029822, 0.029947, 0.030317],
+                        [0.029947, 0.029581, 0.029458, 0.029581, 0.029947],
+                        [0.029822, 0.029458, 0.029336, 0.029458, 0.029822],
+                        [0.029947, 0.029581, 0.029458, 0.029581, 0.029947],
+                        [0.030316, 0.029947, 0.029822, 0.029947, 0.030316],
+                    ],
+                    # Frame 4
+                    [
+                        [0.098142, 0.030317, 0.030191, 0.030317, 0.098142],
+                        [0.030317, 0.029947, 0.029822, 0.029947, 0.030316],
+                        [0.030191, 0.029822, 0.029698, 0.029822, 0.030191],
+                        [0.030317, 0.029947, 0.029822, 0.029947, 0.030316],
+                        [0.098142, 0.030317, 0.030191, 0.030316, 0.098142],
+                    ],
+                ]
+            ]
+        ],
+    ],
+]
+
+
+@skip_if_no_cpp_extension
+class BilateralFilterTestCaseCpuPrecise(unittest.TestCase):
+    @parameterized.expand(TEST_CASES)
+    def test_cpu_precise(self, test_case_description, sigmas, input, expected):
+
+        # Params to determine the implementation to test
+        device = torch.device("cpu")
+
+        # Create input tensor and apply filter
+        input_tensor = torch.from_numpy(np.array(input)).to(dtype=torch.double, device=device)
+
+        len_input = len(input_tensor.shape)
+        # C++ extension so far only supports 5-dim inputs.
+        if len_input == 3:
+            input_tensor = input_tensor.unsqueeze(3).unsqueeze(4)
+        elif len_input == 4:
+            input_tensor = input_tensor.unsqueeze(4)
+
+        output = TrainableBilateralFilterFunction.apply(input_tensor, *sigmas).cpu().numpy()
+
+        # Make sure to return tensor of the same shape as the input.
+        if len_input == 3:
+            output = output.squeeze(4).squeeze(3)
+        elif len_input == 4:
+            output = output.squeeze(4)
+
+        # Ensure result are as expected.
+        np.testing.assert_allclose(output, expected, atol=1e-5)
+
+    @parameterized.expand(TEST_CASES)
+    def test_cpu_precise_backwards(self, test_case_description, sigmas, input, expected):
+
+        # Params to determine the implementation to test
+        device = torch.device("cpu")
+
+        # Prepare input tensor
+        input_tensor = torch.from_numpy(np.array(input)).to(dtype=torch.double, device=device)
+        input_tensor.requires_grad = True
+
+        # C++ extension so far only supports 5-dim inputs.
+        len_input = len(input_tensor.shape)
+        if len_input == 3:
+            input_tensor = input_tensor.unsqueeze(3).unsqueeze(4)
+        elif len_input == 4:
+            input_tensor = input_tensor.unsqueeze(4)
+
+        # Check gradient toward input.
+        args = (
+            input_tensor,
+            torch.tensor(sigmas[0]),
+            torch.tensor(sigmas[1]),
+            torch.tensor(sigmas[2]),
+            torch.tensor(sigmas[3]),
+        )
+        gradcheck(TrainableBilateralFilterFunction.apply, args, eps=1e-6, atol=1e-5, raise_exception=False)
+        input_tensor = input_tensor.detach()
+        input_tensor.requires_grad = False
+
+        # Check gradient toward sigma_x.
+        args = (
+            input_tensor,
+            torch.tensor(sigmas[0], dtype=torch.double, requires_grad=True),
+            torch.tensor(sigmas[1]),
+            torch.tensor(sigmas[2]),
+            torch.tensor(sigmas[3]),
+        )
+        gradcheck(TrainableBilateralFilterFunction.apply, args, eps=1e-2, atol=1e-3, raise_exception=False)
+
+        # Check gradient toward sigma_y.
+        args = (
+            input_tensor,
+            torch.tensor(sigmas[0]),
+            torch.tensor(sigmas[1], dtype=torch.double, requires_grad=True),
+            torch.tensor(sigmas[2]),
+            torch.tensor(sigmas[3]),
+        )
+        gradcheck(TrainableBilateralFilterFunction.apply, args, eps=1e-2, atol=1e-3, raise_exception=False)
+
+        # Check gradient toward sigma_z.
+        args = (
+            input_tensor,
+            torch.tensor(sigmas[0]),
+            torch.tensor(sigmas[1]),
+            torch.tensor(sigmas[2], dtype=torch.double, requires_grad=True),
+            torch.tensor(sigmas[3]),
+        )
+        gradcheck(TrainableBilateralFilterFunction.apply, args, eps=1e-2, atol=1e-3, raise_exception=False)
+
+        # Check gradient toward sigma_color.
+        args = (
+            input_tensor,
+            torch.tensor(sigmas[0]),
+            torch.tensor(sigmas[1]),
+            torch.tensor(sigmas[2]),
+            torch.tensor(sigmas[3], dtype=torch.double, requires_grad=True),
+        )
+        gradcheck(TrainableBilateralFilterFunction.apply, args, eps=1e-3, atol=1e-3, raise_exception=False)
+
+
+@skip_if_no_cuda
+@skip_if_no_cpp_extension
+class BilateralFilterTestCaseCudaPrecise(unittest.TestCase):
+    @parameterized.expand(TEST_CASES)
+    def test_cuda_precise(self, test_case_description, sigmas, input, expected):
+
+        # Skip this test
+        if not torch.cuda.is_available():
+            return
+
+        # Params to determine the implementation to test
+        device = torch.device("cuda")
+
+        # Create input tensor and apply filter
+        input_tensor = torch.from_numpy(np.array(input)).to(dtype=torch.double, device=device)
+
+        len_input = len(input_tensor.shape)
+        # C++ extension so far only supports 5-dim inputs.
+        if len_input == 3:
+            input_tensor = input_tensor.unsqueeze(3).unsqueeze(4)
+        elif len_input == 4:
+            input_tensor = input_tensor.unsqueeze(4)
+
+        output = TrainableBilateralFilterFunction.apply(input_tensor, *sigmas).cpu().numpy()
+
+        # Make sure to return tensor of the same shape as the input.
+        if len_input == 3:
+            output = output.squeeze(4).squeeze(3)
+        elif len_input == 4:
+            output = output.squeeze(4)
+
+        # Ensure result are as expected.
+        np.testing.assert_allclose(output, expected, atol=1e-5)
+
+    @parameterized.expand(TEST_CASES)
+    def test_cuda_precise_backwards(self, test_case_description, sigmas, input, expected):
+
+        # Params to determine the implementation to test
+        device = torch.device("cuda")
+
+        # Prepare input tensor
+        input_tensor = torch.from_numpy(np.array(input)).to(dtype=torch.double, device=device)
+        input_tensor.requires_grad = True
+
+        # C++ extension so far only supports 5-dim inputs.
+        len_input = len(input_tensor.shape)
+        if len_input == 3:
+            input_tensor = input_tensor.unsqueeze(3).unsqueeze(4)
+        elif len_input == 4:
+            input_tensor = input_tensor.unsqueeze(4)
+
+        # Check gradient toward input.
+        args = (
+            input_tensor,
+            torch.tensor(sigmas[0]),
+            torch.tensor(sigmas[1]),
+            torch.tensor(sigmas[2]),
+            torch.tensor(sigmas[3]),
+        )
+        gradcheck(TrainableBilateralFilterFunction.apply, args, eps=1e-6, atol=1e-5, raise_exception=False)
+        input_tensor = input_tensor.detach()
+        input_tensor.requires_grad = False
+
+        # Check gradient toward sigma_x.
+        args = (
+            input_tensor,
+            torch.tensor(sigmas[0], dtype=torch.double, requires_grad=True),
+            torch.tensor(sigmas[1]),
+            torch.tensor(sigmas[2]),
+            torch.tensor(sigmas[3]),
+        )
+        gradcheck(TrainableBilateralFilterFunction.apply, args, eps=1e-2, atol=1e-3, raise_exception=False)
+
+        # Check gradient toward sigma_y.
+        args = (
+            input_tensor,
+            torch.tensor(sigmas[0]),
+            torch.tensor(sigmas[1], dtype=torch.double, requires_grad=True),
+            torch.tensor(sigmas[2]),
+            torch.tensor(sigmas[3]),
+        )
+        gradcheck(TrainableBilateralFilterFunction.apply, args, eps=1e-2, atol=1e-3, raise_exception=False)
+
+        # Check gradient toward sigma_z.
+        args = (
+            input_tensor,
+            torch.tensor(sigmas[0]),
+            torch.tensor(sigmas[1]),
+            torch.tensor(sigmas[2], dtype=torch.double, requires_grad=True),
+            torch.tensor(sigmas[3]),
+        )
+        gradcheck(TrainableBilateralFilterFunction.apply, args, eps=1e-2, atol=1e-3, raise_exception=False)
+
+        # Check gradient toward sigma_color.
+        args = (
+            input_tensor,
+            torch.tensor(sigmas[0]),
+            torch.tensor(sigmas[1]),
+            torch.tensor(sigmas[2]),
+            torch.tensor(sigmas[3], dtype=torch.double, requires_grad=True),
+        )
+        gradcheck(TrainableBilateralFilterFunction.apply, args, eps=1e-3, atol=1e-3, raise_exception=False)
+
+
+if __name__ == "__main__":
+    unittest.main()