initial commit

Author: clee-ai

conf/dataset/default.yaml

@@ -0,0 +1,6 @@
+# @package dataset
+_target_: lightning_transformers.core.data.TransformerDataModule
+cfg:
+  # torch data-loader specific arguments
+  batch_size: ${training.batch_size}
+  num_workers: ${training.num_workers}

conf/dataset/segmentation/default.yaml

@@ -0,0 +1,3 @@
+# @package dataset
+defaults:
+  - /dataset/default

conf/dataset/segmentation/s3dis/s3dis1x1.yaml

@@ -0,0 +1,8 @@
+# @package dataset
+defaults:
+  - segmentation/default
+_target_: lightning_transformers.task.nlp.multiple_choice.RaceMultipleChoiceDataModule
+cfg:
+  dataset_name: race
+  dataset_config_name: 'all'
+  padding: False

lightning_transformers/core/instantiator.py

@@ -0,0 +1,103 @@
+import logging
+from typing import Optional, TYPE_CHECKING, Union
+
+import hydra
+import pytorch_lightning as pl
+import torch
+from omegaconf import DictConfig
+
+from lightning_transformers.core import TransformerDataModule
+from lightning_transformers.core.data import TokenizerDataModule
+
+if TYPE_CHECKING:
+    # avoid circular imports
+    from lightning_transformers.core import TaskTransformer
+
+
+class Instantiator:
+
+    def model(self, *args, **kwargs):
+        raise NotImplementedError("Child class must implement method")
+
+    def optimizer(self, *args, **kwargs):
+        raise NotImplementedError("Child class must implement method")
+
+    def scheduler(self, *args, **kwargs):
+        raise NotImplementedError("Child class must implement method")
+
+    def data_module(self, *args, **kwargs):
+        raise NotImplementedError("Child class must implement method")
+
+    def logger(self, *args, **kwargs):
+        raise NotImplementedError("Child class must implement method")
+
+    def trainer(self, *args, **kwargs):
+        raise NotImplementedError("Child class must implement method")
+
+    def instantiate(self, *args, **kwargs):
+        raise NotImplementedError("Child class must implement method")
+
+
+class HydraInstantiator(Instantiator):
+
+    def model(
+        self,
+        cfg: DictConfig,
+        model_data_kwargs: Optional[DictConfig] = None,
+        tokenizer: Optional[DictConfig] = None,
+        pipeline_kwargs: Optional[DictConfig] = None
+    ) -> "TaskTransformer":
+        if model_data_kwargs is None:
+            model_data_kwargs = {}
+        model_data_kwargs = dict(model_data_kwargs)  # avoid ConfigKeyError: Key 'tokenizer' is not in struct`
+
+        # use `model_data_kwargs` to pass `tokenizer` and `pipeline_kwargs`
+        # as not all models might contain these parameters.
+        if tokenizer:
+            model_data_kwargs["tokenizer"] = self.instantiate(tokenizer)
+        if pipeline_kwargs:
+            model_data_kwargs["pipeline_kwargs"] = pipeline_kwargs
+
+        return self.instantiate(cfg, instantiator=self, **model_data_kwargs)
+
+    def optimizer(self, model: torch.nn.Module, cfg: DictConfig) -> torch.optim.Optimizer:
+        no_decay = ["bias", "LayerNorm.weight"]
+        grouped_parameters = [
+            {
+                "params": [
+                    p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay) and p.requires_grad
+                ],
+                "weight_decay": cfg.weight_decay,
+            },
+            {
+                "params": [
+                    p for n, p in model.named_parameters() if any(nd in n for nd in no_decay) and p.requires_grad
+                ],
+                "weight_decay": 0.0,
+            },
+        ]
+        return self.instantiate(cfg, grouped_parameters)
+
+    def scheduler(self, cfg: DictConfig, optimizer: torch.optim.Optimizer) -> torch.optim.lr_scheduler._LRScheduler:
+        return self.instantiate(cfg, optimizer=optimizer)
+
+    def data_module(
+        self,
+        cfg: DictConfig,
+        tokenizer: Optional[DictConfig] = None,
+    ) -> Union[TransformerDataModule, TokenizerDataModule]:
+        if tokenizer:
+            return self.instantiate(cfg, tokenizer=self.instantiate(tokenizer))
+        return self.instantiate(cfg)
+
+    def logger(self, cfg: DictConfig) -> Optional[logging.Logger]:
+        if cfg.get("log"):
+            if isinstance(cfg.trainer.logger, bool):
+                return cfg.trainer.logger
+            return self.instantiate(cfg.trainer.logger)
+
+    def trainer(self, cfg: DictConfig, **kwargs) -> pl.Trainer:
+        return self.instantiate(cfg, **kwargs)
+
+    def instantiate(self, *args, **kwargs):
+        return hydra.utils.instantiate(*args, **kwargs)

lightning_transformers/dataset/base_dataset.py

@@ -0,0 +1,55 @@
+from typing import Any, Callable, Dict, Optional
+
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader
+
+from lightning_transformers.core.config import TransformerDataConfig
+
+class TransformerDataModule(pl.LightningDataModule):
+
+    def __init__(self, cfg: TransformerDataConfig = TransformerDataConfig()) -> None:
+        super().__init__()
+        self.cfg = cfg
+        self.ds = None
+
+    def train_dataloader(self) -> DataLoader:
+        return DataLoader(
+            self.ds["train"],
+            batch_size=self.batch_size,
+            num_workers=self.cfg.num_workers,
+            collate_fn=self.collate_fn,
+        )
+
+    def val_dataloader(self) -> DataLoader:
+        return DataLoader(
+            self.ds["validation"],
+            batch_size=self.batch_size,
+            num_workers=self.cfg.num_workers,
+            collate_fn=self.collate_fn,
+        )
+
+    def test_dataloader(self) -> Optional[DataLoader]:
+        if "test" in self.ds:
+            return DataLoader(
+                self.ds["test"],
+                batch_size=self.batch_size,
+                num_workers=self.cfg.num_workers,
+                collate_fn=self.collate_fn,
+            )
+
+    @property
+    def batch_size(self) -> int:
+        return self.cfg.batch_size
+
+    @property
+    def collate_fn(self) -> Optional[Callable]:
+        return None
+
+    @property
+    def model_data_kwargs(self) -> Dict:
+        """
+        Override to provide the model with additional kwargs.
+        This is useful to provide the number of classes/pixels to the model or any other data specific args
+        Returns: Dict of args
+        """
+        return {}

lightning_transformers/dataset/geometric_dataset.py

@@ -0,0 +1,78 @@
+
+
+class BaseDataset(TransformerDataModule):
+
+    NAME = ...
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+
+        self._threshold = kwargs.get("threshold", None)
+        self.__instantiate_transform(kwargs)
+        self.clean_kwargs(kwargs)
+        TransformerDataModule.__init__(self, *args, **kwargs)
+
+        self.dataset_train = None
+        self.dataset_val = None
+        self.dataset_test = None
+
+        self._seed = 42
+        self._num_workers = 2
+        self._shuffle = True
+        self._drop_last = False
+        self._pin_memory = True
+        self._follow_batch = []
+
+        self._hyper_parameters = {}
+
+    def __handle_mixin(self):
+        pass
+
+    def clean_kwargs(self, kwargs):
+        LightningDataModuleArgs = inspect.getargspec(LightningDataModule.__init__).args
+        keys = list(kwargs.keys())
+        for key in keys:
+            if key not in LightningDataModuleArgs:
+                del_attr(kwargs, key)
+
+    @property
+    def config(self):
+        return {"dataset_config": {}}
+
+    def __instantiate_transform(self, kwargs):
+        self._pre_transform = None
+        self._transform = None
+        self._train_transform = None
+        self._val_transform = None
+        self._test_transform = None
+
+        for k in [k for k in kwargs]:
+            if "transform" in k and kwargs.get(k) is not None:
+                transforms = []
+                for t in kwargs.get(k):
+                    if t.get("activate") is not None:
+                        if t.activate is False:
+                            continue
+                        del t["activate"]
+                    transforms.append(instantiate(t))
+                transform = T.Compose(transforms)
+                setattr(self, f"_{k}", transform)
+                del kwargs[k]
+
+    @property
+    def num_features(self):
+        pass
+
+    @property
+    def num_classes(self):
+        pass
+
+    @property
+    def hyper_parameters(self):
+        return {"num_features": self.num_features, "num_classes": self.num_classes}
+
+    def prepare_data(self):
+        pass

lightning_transformers/trainer.py

@@ -0,0 +1,75 @@
+from typing import Any, Optional
+
+import hydra
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import LightningDataModule
+from pytorch_lightning.utilities.distributed import rank_zero_info
+
+from lightning_transformers.core import TaskTransformer, TransformerDataModule
+from lightning_transformers.core.config import TaskConfig, TrainerConfig, TransformerDataConfig
+from lightning_transformers.core.instantiator import HydraInstantiator, Instantiator
+from lightning_transformers.core.nlp.config import HFTokenizerConfig
+from lightning_transformers.core.utils import set_ignore_warnings
+
+
+def run(
+    instantiator: Instantiator,
+    ignore_warnings: bool = True,
+    run_test_after_fit: bool = True,
+    dataset: TransformerDataConfig = TransformerDataConfig(),
+    task: TaskConfig = TaskConfig(),
+    trainer: TrainerConfig = TrainerConfig(),
+    tokenizer: Optional[HFTokenizerConfig] = None,
+    logger: Optional[Any] = None,
+) -> None:
+    if ignore_warnings:
+        set_ignore_warnings()
+
+    data_module_kwargs = {}
+    if tokenizer is not None:
+        data_module_kwargs["tokenizer"] = tokenizer
+
+    data_module: TransformerDataModule = instantiator.data_module(dataset, **data_module_kwargs)
+    if data_module is None:
+        raise ValueError("No dataset found. Hydra hint: did you set `dataset=...`?")
+    if not isinstance(data_module, LightningDataModule):
+        raise ValueError(
+            "The instantiator did not return a DataModule instance."
+            " Hydra hint: is `dataset._target_` defined?`"
+        )
+    data_module.setup("fit")
+
+    model: TaskTransformer = instantiator.model(task, model_data_kwargs=getattr(data_module, "model_data_kwargs", None))
+    trainer = instantiator.trainer(
+        trainer,
+        logger=logger,
+    )
+
+    trainer.fit(model, datamodule=data_module)
+    if run_test_after_fit:
+        trainer.test(model, datamodule=data_module)
+
+
+def main(cfg: DictConfig) -> None:
+    rank_zero_info(OmegaConf.to_yaml(cfg))
+    instantiator = HydraInstantiator()
+    logger = instantiator.logger(cfg)
+    run(
+        instantiator,
+        ignore_warnings=cfg.get("ignore_warnings"),
+        run_test_after_fit=cfg.get("training").get("run_test_after_fit"),
+        dataset=cfg.get("dataset"),
+        tokenizer=cfg.get("tokenizer"),
+        task=cfg.get("task"),
+        trainer=cfg.get("trainer"),
+        logger=logger,
+    )
+
+
+@hydra.main(config_path="../../conf", config_name="config")
+def hydra_entry(cfg: DictConfig) -> None:
+    main(cfg)
+
+
+if __name__ == "__main__":
+    hydra_entry()

train.py

@@ -0,0 +1,17 @@
+import hydra
+from hydra.core.global_hydra import GlobalHydra
+from omegaconf import OmegaConf
+from torch_points3d.trainer import Trainer
+
+OmegaConf.register_new_resolver("get_filename", lambda x: x.split('/')[-1])
+@hydra.main(config_path="conf", config_name="config")
+def main(cfg):
+    OmegaConf.set_struct(cfg, False)  # This allows getattr and hasattr methods to function correctly
+    if cfg.pretty_print:
+        print(OmegaConf.to_yaml(cfg))
+
+    trainer = Trainer(cfg)
+    trainer.train()
+
+if __name__ == "__main__":
+    main()