1840 add tbf layer

[faebstn96]

Fixes #1840 .

Description

I integrated the trainable bilateral filter layer (TBF) in the MONAI repository as a new PyTorch filter layer. The TBF contains an analytical gradient derivation toward its filter parameters and its noisy input image which enables gradient-based optimization within the PyTorch graph. See here for more details on the gradient derivation. Unit tests were added that check the filter output as well as the gradient computation.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[wyli]

/black

[wyli]

/black
to compile and run the tests of this PR (assuming that a recent version of pytorch has been installed properly):

git clone https://github.com/faebstn96/MONAI.git --branch 1840-add_tbf_layer
cd MONAI/
pip install -r requirements-min.txt
BUILD_MONAI=1 python setup.py develop
BUILD_MONAI=1 python -m tests.test_trainable_bilateral

[zehuanw]

Hi @faebstn96 , I read your code in the past two weeks and it looks great! Just left two comments directly on the code.

One more question: It seems that most of the logics in the backprop is the same to forward (gaussianKernel_xyz are also the same?). Does it make sense to generate filter_kernel_back (in bf_layer_gpu_backward.cu) directly in forward and compute gradientOutputTensor with gradientInputTensor && the filter? If feasible, backward code will be much simpler.

[faebstn96]

Hi @zehuanw,
Thank you for looking into the code in such detail!

One more question: It seems that most of the logics in the backprop is the same to forward (gaussianKernel_xyz are also the same?). Does it make sense to generate filter_kernel_back (in bf_layer_gpu_backward.cu) directly in forward and compute gradientOutputTensor with gradientInputTensor && the filter? If feasible, backward code will be much simpler.

You are right, the calculation of the spatial Kernels gaussianKernel_xyz is the same in the forward and backward and could be calculated in the forward only and then reused in the backward. However, the intensity range kernel needs to be calculated on the fly as the input intensities change locally (this is the more complex part of the calculation within the CUDA kernels).
I don't think that removing the spatial kernel computation from the backward pass is a change that would notably improve the performance of the code. Usually, the filter size is small in comparison to the filtered image and the spatial filter computation is only done once per filter backward. Please let me know what you think. :)

[wyli]

/build

[wyli]

thanks @faebstn96 and the comments from @zehuanw, the PR looks nice and I'm merging it soon. the fact that the variable notations are consistent with the descriptions in the paper makes the code logic easy to follow...

[wyli]

/build

[faebstn96]

Thank you both @wyli and @zehuanw for your help on this pull request!