Merge pull request #6 from torch-points3d/model

Model

Author: humanpose1

test/test_model.py

@@ -0,0 +1,35 @@
+import unittest
+import sys
+import os
+import torch
+from omegaconf import OmegaConf
+
+from torch_geometric.data import Batch
+
+DIR = os.path.dirname(os.path.realpath(__file__))
+ROOT = os.path.join(DIR, "..")
+sys.path.insert(0, ROOT)
+
+from torch_points3d.models.segmentation.sparseconv3d import APIModel
+
+
+class TestAPIModel(unittest.TestCase):
+    def test_forward(self):
+        option_dataset = OmegaConf.create({"feature_dimension": 1, "num_classes": 10})
+
+        option = OmegaConf.load(os.path.join(ROOT, "conf", "models", "segmentation", "sparseconv3d.yaml"))
+        name_model = list(option.keys())[0]
+        model = APIModel(option[name_model], option_dataset)
+
+        pos = torch.randn(1000, 3)
+        coords = torch.round(pos * 10000)
+        x = torch.ones(1000, 1)
+        batch = torch.zeros(1000).long()
+        y = torch.randint(0, 10, (1000,))
+        data = Batch(pos=pos, x=x, batch=batch, y=y, coords=coords)
+        model.set_input(data)
+        model.forward()
+
+
+if __name__ == "__main__":
+    unittest.main()

torch_points3d/applications/__init__.py

@@ -0,0 +1,99 @@
+from enum import Enum
+from omegaconf import DictConfig
+import logging
+
+from torch_points3d.utils.model_building_utils.model_definition_resolver import resolve
+
+log = logging.getLogger(__name__)
+
+
+class ModelArchitectures(Enum):
+    UNET = "unet"
+    ENCODER = "encoder"
+    DECODER = "decoder"
+
+
+class ModelFactory:
+    MODEL_ARCHITECTURES = [e.value for e in ModelArchitectures]
+
+    @staticmethod
+    def raise_enum_error(arg_name, arg_value, options):
+        raise Exception("The provided argument {} with value {} isn't within {}".format(arg_name, arg_value, options))
+
+    def __init__(
+        self,
+        architecture: str = None,
+        input_nc: int = None,
+        num_layers: int = None,
+        config: DictConfig = None,
+        **kwargs
+    ):
+        if not architecture:
+            raise ValueError()
+        self._architecture = architecture.lower()
+        assert self._architecture in self.MODEL_ARCHITECTURES, ModelFactory.raise_enum_error(
+            "model_architecture", self._architecture, self.MODEL_ARCHITECTURES
+        )
+
+        self._input_nc = input_nc
+        self._num_layers = num_layers
+        self._config = config
+        self._kwargs = kwargs
+
+        if self._config:
+            log.info("The config will be used to build the model")
+
+    @property
+    def modules_lib(self):
+        raise NotImplementedError
+
+    @property
+    def kwargs(self):
+        return self._kwargs
+
+    @property
+    def num_layers(self):
+        return self._num_layers
+
+    @property
+    def num_features(self):
+        return self._input_nc
+
+    def _build_unet(self):
+        raise NotImplementedError
+
+    def _build_encoder(self):
+        raise NotImplementedError
+
+    def _build_decoder(self):
+        raise NotImplementedError
+
+    def build(self):
+        if self._architecture == ModelArchitectures.UNET.value:
+            return self._build_unet()
+        elif self._architecture == ModelArchitectures.ENCODER.value:
+            return self._build_encoder()
+        elif self._architecture == ModelArchitectures.DECODER.value:
+            return self._build_decoder()
+        else:
+            raise NotImplementedError
+
+    @staticmethod
+    def resolve_model(model_config, num_features, kwargs):
+        """Parses the model config and evaluates any expression that may contain constants
+        Overrides any argument in the `define_constants` with keywords wrgument to the constructor
+        """
+        # placeholders to subsitute
+        constants = {
+            "FEAT": max(num_features, 0),
+        }
+
+        # user defined contants to subsitute
+        if "define_constants" in model_config.keys():
+            constants.update(dict(model_config.define_constants))
+            define_constants = model_config.define_constants
+            for key in define_constants.keys():
+                value = kwargs.get(key)
+                if value:
+                    constants[key] = value
+        resolve(model_config, constants)

torch_points3d/applications/modelfactory.py

@@ -0,0 +1,212 @@
+import os
+import sys
+from omegaconf import DictConfig, OmegaConf
+import logging
+import torch
+from torch_geometric.data import Batch
+
+from torch_points3d.applications.modelfactory import ModelFactory
+import torch_points3d.modules.SparseConv3d as sp3d
+from torch_points3d.modules.SparseConv3d.modules import *
+
+# from torch_points3d.core.base_conv.message_passing import *
+# from torch_points3d.core.base_conv.partial_dense import *
+from torch_points3d.models.base_architectures.unet import UnwrappedUnetBasedModel
+from torch_points3d.core.common_modules.base_modules import MLP
+
+from .utils import extract_output_nc
+
+
+CUR_FILE = os.path.realpath(__file__)
+DIR_PATH = os.path.dirname(os.path.realpath(__file__))
+PATH_TO_CONFIG = os.path.join(DIR_PATH, "conf/sparseconv3d")
+
+log = logging.getLogger(__name__)
+
+
+def SparseConv3d(
+    architecture: str = None,
+    input_nc: int = None,
+    num_layers: int = None,
+    config: DictConfig = None,
+    backend: str = "minkowski",
+    *args,
+    **kwargs
+):
+    """Create a Sparse Conv backbone model based on architecture proposed in
+     https://arxiv.org/abs/1904.08755
+
+     Two backends are available at the moment:
+         - https://github.com/mit-han-lab/torchsparse
+         - https://github.com/NVIDIA/MinkowskiEngine
+
+     Parameters
+     ----------
+     architecture : str, optional
+         Architecture of the model, choose from unet, encoder and decoder
+     input_nc : int, optional
+         Number of channels for the input
+    output_nc : int, optional
+         If specified, then we add a fully connected head at the end of the network to provide the requested dimension
+     num_layers : int, optional
+         Depth of the network
+     config : DictConfig, optional
+         Custom config, overrides the num_layers and architecture parameters
+     block:
+         Type of resnet block, ResBlock by default but can be any of the blocks in modules/SparseConv3d/modules.py
+     backend:
+         torchsparse or minkowski
+    """
+    if "SPARSE_BACKEND" in os.environ and sp3d.nn.backend_valid(os.environ["SPARSE_BACKEND"]):
+        sp3d.nn.set_backend(os.environ["SPARSE_BACKEND"])
+    else:
+        sp3d.nn.set_backend(backend)
+
+    factory = SparseConv3dFactory(
+        architecture=architecture, num_layers=num_layers, input_nc=input_nc, config=config, **kwargs
+    )
+    return factory.build()
+
+
+class SparseConv3dFactory(ModelFactory):
+    def _build_unet(self):
+        if self._config:
+            model_config = self._config
+        else:
+            path_to_model = os.path.join(PATH_TO_CONFIG, "unet_{}.yaml".format(self.num_layers))
+            model_config = OmegaConf.load(path_to_model)
+        ModelFactory.resolve_model(model_config, self.num_features, self._kwargs)
+        modules_lib = sys.modules[__name__]
+        return SparseConv3dUnet(model_config, None, modules_lib, **self.kwargs)
+
+    def _build_encoder(self):
+        if self._config:
+            model_config = self._config
+        else:
+            path_to_model = os.path.join(
+                PATH_TO_CONFIG,
+                "encoder_{}.yaml".format(self.num_layers),
+            )
+            model_config = OmegaConf.load(path_to_model)
+        ModelFactory.resolve_model(model_config, self.num_features, self._kwargs)
+        modules_lib = sys.modules[__name__]
+        return SparseConv3dEncoder(model_config, None, modules_lib, **self.kwargs)
+
+
+class BaseSparseConv3d(UnwrappedUnetBasedModel):
+    CONV_TYPE = "sparse"
+
+    def __init__(self, model_config, model_type, modules, *args, **kwargs):
+        super().__init__(model_config, model_type, modules)
+        self.weight_initialization()
+        default_output_nc = kwargs.get("default_output_nc", None)
+        if not default_output_nc:
+            default_output_nc = extract_output_nc(model_config)
+
+        self._output_nc = default_output_nc
+        self._has_mlp_head = False
+        if "output_nc" in kwargs:
+            self._has_mlp_head = True
+            self._output_nc = kwargs["output_nc"]
+            self.mlp = MLP([default_output_nc, self.output_nc], activation=torch.nn.ReLU(), bias=False)
+
+    @property
+    def has_mlp_head(self):
+        return self._has_mlp_head
+
+    @property
+    def output_nc(self):
+        return self._output_nc
+
+    def weight_initialization(self):
+        for m in self.modules():
+            if isinstance(m, sp3d.nn.Conv3d) or isinstance(m, sp3d.nn.Conv3dTranspose):
+                torch.nn.init.kaiming_normal_(m.kernel, mode="fan_out", nonlinearity="relu")
+
+            if isinstance(m, sp3d.nn.BatchNorm):
+                torch.nn.init.constant_(m.bn.weight, 1)
+                torch.nn.init.constant_(m.bn.bias, 0)
+
+    def _set_input(self, data):
+        """Unpack input data from the dataloader and perform necessary pre-processing steps.
+
+        Parameters
+        -----------
+        data:
+            a dictionary that contains the data itself and its metadata information.
+        """
+        self.input = sp3d.nn.SparseTensor(data.x, data.coords, data.batch)
+        if data.pos is not None:
+            self.xyz = data.pos
+        else:
+            self.xyz = data.coords
+
+
+class SparseConv3dEncoder(BaseSparseConv3d):
+    def forward(self, data, *args, **kwargs):
+        """
+        Parameters:
+        -----------
+        data
+            A SparseTensor that contains the data itself and its metadata information. Should contain
+                F -- Features [N, C]
+                coords -- Coords [N, 4]
+
+        Returns
+        --------
+        data:
+            - x [1, output_nc]
+
+        """
+        self._set_input(data)
+        data = self.input
+        for i in range(len(self.down_modules)):
+            data = self.down_modules[i](data)
+
+        out = Batch(x=data.F, batch=data.C[:, 0].long().to(data.F.device))
+        if not isinstance(self.inner_modules[0], Identity):
+            out = self.inner_modules[0](out)
+
+        if self.has_mlp_head:
+            out.x = self.mlp(out.x)
+        return out
+
+
+class SparseConv3dUnet(BaseSparseConv3d):
+    def forward(self, data, *args, **kwargs):
+        """Run forward pass.
+        Input --- D1 -- D2 -- D3 -- U1 -- U2 -- output
+                   |      |_________|     |
+                   |______________________|
+
+        Parameters
+        -----------
+        data
+            A SparseTensor that contains the data itself and its metadata information. Should contain
+                F -- Features [N, C]
+                coords -- Coords [N, 4]
+
+        Returns
+        --------
+        data:
+            - pos [N, 3] (coords or real pos if xyz is in data)
+            - x [N, output_nc]
+            - batch [N]
+        """
+        self._set_input(data)
+        data = self.input
+        stack_down = []
+        for i in range(len(self.down_modules) - 1):
+            data = self.down_modules[i](data)
+            stack_down.append(data)
+
+        data = self.down_modules[-1](data)
+        stack_down.append(None)
+        # TODO : Manage the inner module
+        for i in range(len(self.up_modules)):
+            data = self.up_modules[i](data, stack_down.pop())
+
+        out = Batch(x=data.F, pos=self.xyz)
+        if self.has_mlp_head:
+            out.x = self.mlp(out.x)
+        return out

torch_points3d/applications/sparseconv3d.py

@@ -0,0 +1,9 @@
+def extract_output_nc(model_config):
+    """Extracts the number of channels at the output of the network form the model config"""
+    if model_config.get("up_conv") is not None:
+        output_nc = model_config.up_conv.up_conv_nn[-1][-1]
+    elif model_config.get("innermost") is not None:
+        output_nc = model_config.innermost.nn[-1]
+    else:
+        raise ValueError("Input model_config does not match expected pattern")
+    return output_nc

torch_points3d/applications/utils.py

@@ -0,0 +1,2 @@
+from .base_modules import *
+from .spatial_transform import *