Repository: Project-MONAI/MONAI

Issue:
Add sample CUDA program based on PyTorch C++ extension
**Is your feature request related to a problem? Please describe.**
In order to set up CUDA CI / CD support, we need to add some test program for CUDA extension first.
May replace with some other programs when we have enough CUDA kernels later.



Discussion:

Resolving patch:
diff --git a/monai/__init__.py b/monai/__init__.py
--- a/monai/__init__.py
+++ b/monai/__init__.py
@@ -22,7 +22,9 @@
 
 __basedir__ = os.path.dirname(__file__)
 
-excludes = "^(handlers)"  # the handlers have some external decorators the users may not have installed
+# handlers_* have some external decorators the users may not have installed
+# *.so files and folder "_C" may not exist when the cpp extensions are not compiled
+excludes = "(^(handlers))|((\\.so)$)|(_C)"
 
 # load directory modules only, skip loading individual files
 load_submodules(sys.modules[__name__], False, exclude_pattern=excludes)
diff --git a/monai/networks/layers/simplelayers.py b/monai/networks/layers/simplelayers.py
--- a/monai/networks/layers/simplelayers.py
+++ b/monai/networks/layers/simplelayers.py
@@ -10,15 +10,19 @@
 # limitations under the License.
 
 from typing import Sequence, Union
-
+import math
 import torch
-import torch.nn as nn
+from torch import nn
+from torch.autograd import Function
 import torch.nn.functional as F
-
 from monai.networks.layers.convutils import gaussian_1d, same_padding
 from monai.utils import ensure_tuple_rep
+from monai.utils import optional_import
 
-__all__ = ["SkipConnection", "Flatten", "GaussianFilter"]
+_C, _ = optional_import("monai._C")
+_C_CUDA, _ = optional_import("monai._C_CUDA")
+
+__all__ = ["SkipConnection", "Flatten", "GaussianFilter", "LLTM"]
 
 
 class SkipConnection(nn.Module):
@@ -113,3 +117,58 @@ def _conv(input_: torch.Tensor, d: int) -> torch.Tensor:
             return self.conv_n(input=_conv(input_, d - 1), weight=kernel, padding=padding, groups=chns)
 
         return _conv(x, sp_dim - 1)
+
+
+class LLTMFunction(Function):
+    @staticmethod
+    def forward(ctx, input, weights, bias, old_h, old_cell):
+        ext = _C_CUDA if weights.is_cuda else _C
+        outputs = ext.lltm_forward(input, weights, bias, old_h, old_cell)
+        new_h, new_cell = outputs[:2]
+        variables = outputs[1:] + [weights]
+        ctx.save_for_backward(*variables)
+
+        return new_h, new_cell
+
+    @staticmethod
+    def backward(ctx, grad_h, grad_cell):
+        if grad_h.is_cuda:
+            outputs = _C_CUDA.lltm_backward(grad_h.contiguous(), grad_cell.contiguous(), *ctx.saved_tensors)
+            d_old_h, d_input, d_weights, d_bias, d_old_cell, d_gates = outputs
+        else:
+            outputs = _C.lltm_backward(grad_h.contiguous(), grad_cell.contiguous(), *ctx.saved_tensors)
+            d_old_h, d_input, d_weights, d_bias, d_old_cell = outputs
+
+        return d_input, d_weights, d_bias, d_old_h, d_old_cell
+
+
+class LLTM(nn.Module):
+    """
+    This recurrent unit is similar to an LSTM, but differs in that it lacks a forget
+    gate and uses an Exponential Linear Unit (ELU) as its internal activation function.
+    Because this unit never forgets, call it LLTM, or Long-Long-Term-Memory unit.
+    It has both C++ and CUDA implementation, automatically switch according to the
+    target device where put this module to.
+
+    Args:
+        input_features: size of input feature data
+        state_size: size of the state of recurrent unit
+
+    Referring to: https://pytorch.org/tutorials/advanced/cpp_extension.html
+    """
+
+    def __init__(self, input_features: int, state_size: int):
+        super(LLTM, self).__init__()
+        self.input_features = input_features
+        self.state_size = state_size
+        self.weights = nn.Parameter(torch.empty(3 * state_size, input_features + state_size))
+        self.bias = nn.Parameter(torch.empty(1, 3 * state_size))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        stdv = 1.0 / math.sqrt(self.state_size)
+        for weight in self.parameters():
+            weight.data.uniform_(-stdv, +stdv)
+
+    def forward(self, input, state):
+        return LLTMFunction.apply(input, self.weights, self.bias, *state)
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -9,15 +9,51 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
+
 from setuptools import find_packages, setup
 
 import versioneer
 
-if __name__ == "__main__":
-    setup(
-        version=versioneer.get_version(),
-        cmdclass=versioneer.get_cmdclass(),
-        packages=find_packages(exclude=("docs", "examples", "tests", "research")),
-        zip_safe=False,
-        package_data={"monai": ["py.typed"]},
-    )
+
+def get_extensions():
+
+    try:
+        import torch
+        from torch.utils.cpp_extension import CUDA_HOME, CppExtension, CUDAExtension
+
+        print(f"setup.py with torch {torch.__version__}")
+    except ImportError:
+        warnings.warn("torch cpp/cuda building skipped.")
+        return []
+
+    ext_modules = [CppExtension("monai._C", ["monai/networks/extensions/lltm/lltm.cpp"])]
+    if torch.cuda.is_available() and (CUDA_HOME is not None):
+        ext_modules.append(
+            CUDAExtension(
+                "monai._C_CUDA",
+                ["monai/networks/extensions/lltm/lltm_cuda.cpp", "monai/networks/extensions/lltm/lltm_cuda_kernel.cu"],
+            )
+        )
+    return ext_modules
+
+
+def get_cmds():
+    cmds = versioneer.get_cmdclass()
+    try:
+        from torch.utils.cpp_extension import BuildExtension
+
+        cmds.update({"build_ext": BuildExtension})
+    except ImportError:
+        warnings.warn("torch cpp_extension module not found.")
+    return cmds
+
+
+setup(
+    version=versioneer.get_version(),
+    cmdclass=get_cmds(),
+    packages=find_packages(exclude=("docs", "examples", "tests", "research")),
+    zip_safe=False,
+    package_data={"monai": ["py.typed"]},
+    ext_modules=get_extensions(),
+)


Test patch:
diff --git a/tests/test_lltm.py b/tests/test_lltm.py
new file mode 100644
--- /dev/null
+++ b/tests/test_lltm.py
@@ -0,0 +1,55 @@
+# Copyright 2020 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.
+
+import unittest
+
+import torch
+from parameterized import parameterized
+
+from monai.networks.layers import LLTM
+
+TEST_CASE_1 = [
+    {"input_features": 32, "state_size": 2},
+    torch.tensor([[-0.1622, 0.1663], [0.5465, 0.0459], [-0.1436, 0.6171], [0.3632, -0.0111]]),
+    torch.tensor([[-1.3773, 0.3348], [0.8353, 1.3064], [-0.2179, 4.1739], [1.3045, -0.1444]]),
+]
+
+
+class TestLLTM(unittest.TestCase):
+    @parameterized.expand([TEST_CASE_1])
+    def test_value(self, input_param, expected_h, expected_c):
+        torch.manual_seed(0)
+        x = torch.randn(4, 32)
+        h = torch.randn(4, 2)
+        c = torch.randn(4, 2)
+        new_h, new_c = LLTM(**input_param)(x, (h, c))
+        (new_h.sum() + new_c.sum()).backward()
+
+        torch.testing.assert_allclose(new_h, expected_h, rtol=0.0001, atol=1e-04)
+        torch.testing.assert_allclose(new_c, expected_c, rtol=0.0001, atol=1e-04)
+
+    @parameterized.expand([TEST_CASE_1])
+    def test_value_cuda(self, input_param, expected_h, expected_c):
+        device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu:0")
+        torch.manual_seed(0)
+        x = torch.randn(4, 32).to(device)
+        h = torch.randn(4, 2).to(device)
+        c = torch.randn(4, 2).to(device)
+        lltm = LLTM(**input_param).to(device)
+        new_h, new_c = lltm(x, (h, c))
+        (new_h.sum() + new_c.sum()).backward()
+
+        torch.testing.assert_allclose(new_h, expected_h.to(device), rtol=0.0001, atol=1e-04)
+        torch.testing.assert_allclose(new_c, expected_c.to(device), rtol=0.0001, atol=1e-04)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/testing_data/._Task04_Hippocampus b/tests/testing_data/._Task04_Hippocampus
deleted file mode 100755
Binary files a/tests/testing_data/._Task04_Hippocampus and /dev/null differ