Added HoverNet Loss

[JHancox]

Fixes #4672

Description

Adds the Hovernet Loss, which takes the raw predictions from the HoverNet net and produces a Loss Tensor.

Hovernet Loss added along with test script that verfies same results as original TIA version of Hovernet loss. Test script also include synthetic image generation (fixed to deterministic behaviour)

N.B. Currently does not include a test function for torch script since the test includes code that torchscipt cannot compile - mainly the SobelGradient and associated functions. Need to resolve this.

Types of changes

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

[Nic-Ma]

Hi @JHancox ,

I would suggest to put this loss implementation in: https://github.com/Project-MONAI/MONAI/tree/dev/monai/apps/pathology /losses/hovernet_loss.py

What do you think?

Thanks.

[JHancox]

Hi @JHancox ,

I would suggest to put this loss implementation in: https://github.com/Project-MONAI/MONAI/tree/dev/monai/apps/pathology /losses/hovernet_loss.py

What do you think?

Thanks.

Thanks @Nic-Ma - I see your logic and it would make sense. On the other hand HoverNet is already in the monai/networks/nets branch and so we'd probably need to move this too. @drbeh - what do you think?

[Nic-Ma]

Hi @JHancox ,

I think the monai/networks/nets/hovernet is OK. I requested to move the loss to apps/pathology because it seems not a general loss function, it can only work for the particular HoverNet training pipeline. Please feel free to correct me if I misunderstand anything.

Thanks.

[JHancox]

Thanks @Nic-Ma. You are right - the loss is specific to Hovernet. I have no problem redirecting this to apps but will leave the decision to @drbeh

[bhashemian]

@JHancox @Nic-Ma I agree. This is better suited for monai/apps/pathology/loss/.

[bhashemian]

Hi @Nic-Ma @wyli @rijobro @ericspod @yiheng-wang-nv, Do you have any insight why torchscript is failing here? https://github.com/Project-MONAI/MONAI/actions/runs/3190912645/jobs/5206609165#step:7:30985

It seems that it is related to SobelGradients but do you know what's the issue?

Thanks for your help.

ERROR: test_script (tests.test_hovernet_loss.TestHoverNetLoss)
tests finished, printing completed times >10.0s in ascending order...

test_senet_shape_0 (tests.test_senet.TestPretrainedSENET) (10.9s)
test_cpu_precise_backwards_7_3_dimension_1_channel_high_spatial_sigma_high_color_sigma (tests.test_bilateral_precise.BilateralFilterTestCaseCpuPrecise) (11.3s)
test_invert (tests.test_invertd.TestInvertd) (12.7s)
test_script_0 (tests.test_senet.TestSENET) (15.9s)
test_cuda_1_1_batches_1_dimensions_5_channels_2_classes_1_mixtures (tests.test_gmm.GMMTestCase) (16.5s)
test_cuda_2_1_batches_2_dimensions_2_channels_4_classes_4_mixtures (tests.test_gmm.GMMTestCase) (16.7s)
test_cuda_3_1_batches_3_dimensions_1_channels_2_classes_1_mixtures (tests.test_gmm.GMMTestCase) (17.2s)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/__w/MONAI/MONAI/tests/test_hovernet_loss.py", line 181, in test_script
    test_script_save(loss, prediction, target)
  File "/__w/MONAI/MONAI/tests/utils.py", line 704, in test_script_save
    atol=atol,
  File "/__w/MONAI/MONAI/monai/networks/utils.py", line 593, in convert_to_torchscript
    script_module = torch.jit.script(model, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/torch/jit/_script.py", line 1287, in script
    obj, torch.jit._recursive.infer_methods_to_compile
  File "/usr/local/lib/python3.7/dist-packages/torch/jit/_recursive.py", line 458, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/usr/local/lib/python3.7/dist-packages/torch/jit/_recursive.py", line 524, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
  File "/usr/local/lib/python3.7/dist-packages/torch/jit/_recursive.py", line 375, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
  File "/usr/local/lib/python3.7/dist-packages/torch/jit/_recursive.py", line 876, in compile_unbound_method
    create_methods_and_properties_from_stubs(concrete_type, (stub,), ())
  File "/usr/local/lib/python3.7/dist-packages/torch/jit/_recursive.py", line 375, in create_methods_and_properties_from_stubs
test_cuda_0_2_batches_1_dimensions_1_channels_2_classes_2_mixtures (tests.test_gmm.GMMTestCase) (17.6s)
test_dints_inference_3 (tests.test_dints_network.TestDints) (17.7s)
test_dints_inference_1 (tests.test_dints_network.TestDints) (19.0s)
test_dints_inference_2 (tests.test_dints_network.TestDints) (19.7s)
test_run_algo_after_move_files (tests.test_auto3dseg_hpo.TestHPO) (25.9s)
test_run_algo (tests.test_auto3dseg_hpo.TestHPO) (26.0s)
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
  File "/usr/local/lib/python3.7/dist-packages/torch/jit/_recursive.py", line 876, in compile_unbound_method
    create_methods_and_properties_from_stubs(concrete_type, (stub,), ())
  File "/usr/local/lib/python3.7/dist-packages/torch/jit/_recursive.py", line 375, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
RuntimeError: 
Module 'HoVerNetLoss' has no attribute 'sobel' (This attribute exists on the Python module, but we failed to convert Python type: 'monai.transforms.post.array.SobelGradients' to a TorchScript type. Only tensors and (possibly nested) tuples of tensors, lists, or dictsare supported as inputs or outputs of traced functions, but instead got value of type SobelGradients.. Its type was inferred; try adding a type annotation for the attribute.):
  File "/__w/MONAI/MONAI/monai/apps/pathology/losses/hovernet_loss.py", line 56
    
        batch_size = image.shape[0]
        result_h = self.sobel(torch.squeeze(image[:, 0], dim=1))[batch_size:]
                   ~~~~~~~~~~ <--- HERE
        result_v = self.sobel(torch.squeeze(image[:, 1], dim=1))[:batch_size]
    
'HoVerNetLoss._compute_sobel' is being compiled since it was called from 'HoVerNetLoss._mse_gradient_loss'
  File "/__w/MONAI/MONAI/monai/apps/pathology/losses/hovernet_loss.py", line 63
    def _mse_gradient_loss(self, prediction: torch.Tensor, target: torch.Tensor, focus: torch.Tensor) -> torch.Tensor:
    
        pred_grad = self._compute_sobel(prediction)
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        true_grad = self._compute_sobel(target)
    
'HoVerNetLoss._mse_gradient_loss' is being compiled since it was called from 'HoVerNetLoss.forward'
  File "/__w/MONAI/MONAI/monai/apps/pathology/losses/hovernet_loss.py", line 118
        # Use the nuclei class, one hot encoded, as the mask
        loss_hv_mse_grad = (
            self._mse_gradient_loss(
            ~~~~~~~~~~~~~~~~~~~~~~~~
                prediction[HoVerNetBranch.HV.value],
                ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
                target[HoVerNetBranch.HV.value],
                ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
                target[HoVerNetBranch.NP.value][:, 1],
                ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
            )
            * self.lambda_hv_mse_grad


----------------------------------------------------------------------

[Nic-Ma]

Hi @benjamin-gorman ,

As your error message shows: Module 'HoVerNetLoss' has no attribute 'sobel' (This attribute exists on the Python module, but we failed to convert Python type: 'monai.transforms.post.array.SobelGradients' to a TorchScript type. Only tensors and (possibly nested) tuples of tensors, lists, or dictsare supported as inputs or outputs of traced functions, but instead got value of type SobelGradients.. Its type was inferred; try adding a type annotation for the attribute. Maybe you can have a try?

Thanks.

[JHancox]

Hi @benjamin-gorman ,

As your error message shows: Module 'HoVerNetLoss' has no attribute 'sobel' (This attribute exists on the Python module, but we failed to convert Python type: 'monai.transforms.post.array.SobelGradients' to a TorchScript type. Only tensors and (possibly nested) tuples of tensors, lists, or dictsare supported as inputs or outputs of traced functions, but instead got value of type SobelGradients.. Its type was inferred; try adding a type annotation for the attribute. Maybe you can have a try?

Thanks.

I think the problem is that SobelGradients is a Transform, which is a Callable and I don't think TS supports the Callable type. Therefore, perhaps we might refactor SobelGradients so that it can be used as a Tensor-returning class and, perhaps wrap that in the Transform interface. ? Oh, and did you mean @drbeh ? :)

[bhashemian]

As your error message shows: Module 'HoVerNetLoss' has no attribute 'sobel' (This attribute exists on the Python module, but we failed to convert Python type: 'monai.transforms.post.array.SobelGradients' to a TorchScript type. Only tensors and (possibly nested) tuples of tensors, lists, or dictsare supported as inputs or outputs of traced functions, but instead got value of type SobelGradients.. Its type was inferred; try adding a type annotation for the attribute. Maybe you can have a try? Thanks.

I think the problem is that SobelGradients is a Transform, which is a Callable and I don't think TS supports the Callable type. Therefore, perhaps we might refactor SobelGradients so that it can be used as a Tensor-returning class and, perhaps wrap that in the Transform interface. ? Oh, and did you mean @drbeh ? :)

Hi @Nic-Ma, As @JHancox mentioned, we have tried that and didn't work. Even removing sobel from being an attribute and apply it in the call does not work and give other different complaints. In my experience torch script error messages are not that useful although their statement is correct.

@JHancox, this transform is already a Tensor-returning class but the issue with torch script is that it is very limited in the scope of what it can support. https://pytorch.org/docs/stable/jit_unsupported.html#jit-unsupported

@Nic-Ma do we need the loss function to be torch scriptable? I am not sure if it is worth spending more time on this to make it to work with torch script (if possible at all). If it is not required, I would suggest to move forward without it.

[Nic-Ma]

Hi @drbeh ,

Thanks for the quick update. I didn't see any TorchScript use case of loss function, so I feel we can remove the TorchScript tests so far. But I see we have TorchScript compatible tests for almost all the losses: https://github.com/Project-MONAI/MONAI/blob/dev/tests/test_dice_loss.py#L187 @wyli @ericspod Do you know the use case or reason to support that?

Thanks.

[wyli]

agree with @Nic-Ma to remove test_script from TestHoverNetLoss (previously to have torchscript is potentially speed up the training https://github.com/Project-MONAI/MONAI/issues/1592#issuecomment-782182711)

[Nic-Ma]

/build

[Nic-Ma]

Is it ready for review now?

Thanks.

[JHancox]

Is it ready for review now?

Thanks.

I just need to check with @drbeh whether there was anything else. Will try to get an answer today (in Sweden at the moment).

[bhashemian]

I try to resolve some failing tests and some formatting, and then it should be ready. I'll request to review afterwards.

[Nic-Ma]

/build

[Nic-Ma]

/build

[wyli]

this PR causes some integration errors, could you please take a look? https://github.com/Project-MONAI/MONAI/issues/5393 @drbeh @Nic-Ma

[bhashemian]

Hi @wyli, sure, I'll look into it.