Repository: Project-MONAI/MONAI

Issue:
integration TestDynUNetWithInstanceNorm3dNVFuser
**Describe the bug**
logs using pytorch image 22.04 https://github.com/Project-MONAI/MONAI/runs/6306011370?check_suite_focus=true

(| NVIDIA-SMI 495.29.05    Driver Version: 495.29.05    CUDA Version: 11.5   )
```py
ERROR: test_consistency_0 (tests.test_dynunet.TestDynUNetWithInstanceNorm3dNVFuser)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/opt/conda/lib/python3.8/site-packages/parameterized/parameterized.py", line 533, in standalone_func
test_compute_0 (tests.test_prepare_batch_default_dist.DistributedPrepareBatchDefault) (10.5s)
test_compute_1 (tests.test_prepare_batch_default_dist.DistributedPrepareBatchDefault) (10.6s)
test_verify_0__tmp_tmp_HuvzKkaSgJ_tests_testing_data_metadata_json (tests.test_bundle_verify_net.TestVerifyNetwork) (11.2s)
    return func(*(a + p.args), **p.kwargs)
  File "/tmp/tmp.HuvzKkaSgJ/tests/test_dynunet.py", line 146, in test_consistency
    result_fuser = net_fuser(input_tensor)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/monai/networks/nets/dynunet.py", line 268, in forward
    out = self.skip_layers(x)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/monai/networks/nets/dynunet.py", line 46, in forward
    downout = self.downsample(x)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/monai/networks/blocks/dynunet_block.py", line 80, in forward
    out = self.norm1(out)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/apex/normalization/instance_norm.py", line 138, in forward
    out = InstanceNormNVFuserFunction.apply(
  File "/opt/conda/lib/python3.8/site-packages/apex/normalization/instance_norm.py", line 16, in forward
    instance_norm_nvfuser_cuda = importlib.import_module("instance_norm_nvfuser_cuda")
  File "/opt/conda/lib/python3.8/importlib/__init__.py", line 127, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "<frozen importlib._bootstrap>", line 1014, in _gcd_import
  File "<frozen importlib._bootstrap>", line 991, in _find_and_load
  File "<frozen importlib._bootstrap>", line 975, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 657, in _load_unlocked
  File "<frozen importlib._bootstrap>", line 556, in module_from_spec
  File "<frozen importlib._bootstrap_external>", line 1166, in create_module
  File "<frozen importlib._bootstrap>", line 219, in _call_with_frames_removed
ImportError: /opt/conda/lib/python3.8/site-packages/instance_norm_nvfuser_cuda.cpython-38-x86_64-linux-gnu.so: undefined symbol: _ZTVN5torch3jit5fuser4cuda3kir6KernelE

----------------------------------------------------------------------
```


Discussion:
Hi @yiheng-wang-nv ,

Could you please help investigate and fix this `InstanceNorm3dNVFuser` issue in PyTorch docker 22.04?

Thanks in advance.


Resolving patch:
diff --git a/monai/networks/layers/factories.py b/monai/networks/layers/factories.py
--- a/monai/networks/layers/factories.py
+++ b/monai/networks/layers/factories.py
@@ -63,12 +63,14 @@ def use_factory(fact_args):
 import warnings
 from typing import Any, Callable, Dict, Tuple, Type, Union
 
+import torch
 import torch.nn as nn
 
 from monai.utils import look_up_option, optional_import
 
 InstanceNorm3dNVFuser, has_nvfuser = optional_import("apex.normalization", name="InstanceNorm3dNVFuser")
 
+
 __all__ = ["LayerFactory", "Dropout", "Norm", "Act", "Conv", "Pool", "Pad", "split_args"]
 
 
@@ -263,8 +265,21 @@ def instance_nvfuser_factory(dim):
     if dim != 3:
         warnings.warn(f"`InstanceNorm3dNVFuser` only supports 3d cases, use {types[dim - 1]} instead.")
         return types[dim - 1]
-    if not has_nvfuser:
-        warnings.warn("`apex.normalization.InstanceNorm3dNVFuser` is not found, use nn.InstanceNorm3d instead.")
+    # test InstanceNorm3dNVFuser installation with a basic example
+    has_nvfuser_flag = has_nvfuser
+    if not torch.cuda.is_available():
+        return nn.InstanceNorm3d
+    try:
+        layer = InstanceNorm3dNVFuser(num_features=1, affine=True).to("cuda:0")
+        inp = torch.randn([1, 1, 1, 1, 1]).to("cuda:0")
+        out = layer(inp)
+        del inp, out, layer
+    except Exception:
+        has_nvfuser_flag = False
+    if not has_nvfuser_flag:
+        warnings.warn(
+            "`apex.normalization.InstanceNorm3dNVFuser` is not installed properly, use nn.InstanceNorm3d instead."
+        )
         return nn.InstanceNorm3d
     return InstanceNorm3dNVFuser
 


Test patch:
diff --git a/tests/test_dynunet.py b/tests/test_dynunet.py
--- a/tests/test_dynunet.py
+++ b/tests/test_dynunet.py
@@ -17,11 +17,9 @@
 
 from monai.networks import eval_mode
 from monai.networks.nets import DynUNet
-from monai.utils import optional_import
+from monai.utils.module import pytorch_after
 from tests.utils import skip_if_no_cuda, skip_if_windows, test_script_save
 
-_, has_nvfuser = optional_import("apex.normalization", name="InstanceNorm3dNVFuser")
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 strides: Sequence[Union[Sequence[int], int]]
@@ -123,7 +121,6 @@ def test_script(self):
 
 @skip_if_no_cuda
 @skip_if_windows
-@unittest.skipUnless(has_nvfuser, "To use `instance_nvfuser`, `apex.normalization.InstanceNorm3dNVFuser` is needed.")
 class TestDynUNetWithInstanceNorm3dNVFuser(unittest.TestCase):
     @parameterized.expand([TEST_CASE_DYNUNET_3D[0]])
     def test_consistency(self, input_param, input_shape, _):
@@ -145,7 +142,11 @@ def test_consistency(self, input_param, input_shape, _):
                         with eval_mode(net_fuser):
                             result_fuser = net_fuser(input_tensor)
 
-                        torch.testing.assert_close(result, result_fuser)
+                        # torch.testing.assert_allclose() is deprecated since 1.12 and will be removed in 1.14
+                        if pytorch_after(1, 12):
+                            torch.testing.assert_close(result, result_fuser)
+                        else:
+                            torch.testing.assert_allclose(result, result_fuser)
 
 
 class TestDynUNetDeepSupervision(unittest.TestCase):