simple Finetuning (torch-points3d#455)

* add set weights

* set pretrained weight

* add pretrained possibilities

* pre trained seems to work

* small change

* update changelogs

* fix in trainer

* load weights with the same size only

Author: humanpose1

CHANGELOG.md

@@ -10,7 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Support for the IRALab benchmark (https://arxiv.org/abs/2003.12841), with data from the ETH, Canadian Planetary, Kaist and TUM datasets.
-
+- Possibility to load pretrained models by adding the path in the confs for finetuning.
 
 ### Bug fix
 

conf/data/registration/testeth.yaml

@@ -4,8 +4,8 @@ data:
   dataroot: data
   first_subsampling: 0.02
   max_dist_overlap: 0.05
-  min_size_block: 1.5
-  max_size_block: 2
+  min_size_block: 5
+  max_size_block: 7
   num_pos_pairs: 30000
   min_points: 300
   num_points: 5000
@@ -21,10 +21,10 @@ data:
   ss_transform:
     - transform: CubeCrop
       params:
-        c: 1
+        c: 5
     - transform: CubeCrop
       params:
-        c: 1.5
+        c: 5.5
 
   train_transform:
     - transform: SaveOriginalPosId

conf/models/registration/minkowski.yaml

@@ -39,6 +39,48 @@ models:
             num_hn_samples: 256
             pos_thresh: 0.1
             neg_thresh: 1.4
+    MinkUNet_Fragment_pretrained:
+        class: minkowski.MinkowskiFragment
+        path_pretrained: "INSERT PATH"
+        weight_name: "latest"
+        conv_type: "SPARSE"
+        loss_mode: "match"
+        down_conv:
+            module_name: Res2BlockDown
+            dimension: 3
+            bn_momentum: 0.05
+            down_conv_nn:
+                [
+                    [FEAT, 32],
+                    [32, 64],
+                    [64, 128],
+                    [128, 256]
+                ]
+            kernel_size: [5, 3, 3, 3]
+            stride: [1, 2, 2, 2]
+            dilation: [1, 1, 1, 1]
+        up_conv:
+            module_name: Res2BlockUp
+            dimension: 3
+            bn_momentum: 0.05
+            up_conv_nn:
+                [
+                    [256, 64],
+                    [64 + 128, 64],
+                    [64 + 64, 64],
+                    [64 + 32, 64, 32]
+                ]
+            kernel_size: [3, 3, 3, 1]
+            stride: [2, 2, 2, 1]
+            dilation: [1, 1, 1, 1]
+        normalize_feature: True
+        metric_loss:
+          class: "ContrastiveHardestNegativeLoss"
+          params:
+            num_pos: 1024
+            num_hn_samples: 256
+            pos_thresh: 0.1
+            neg_thresh: 1.4
 
 
     Res16UNet32B:

test/test_basemodel.py

@@ -1,4 +1,5 @@
 import unittest
+import torch
 from omegaconf import OmegaConf, DictConfig
 from torch.nn import (
     Sequential,
@@ -49,6 +50,20 @@ def set_input(self, a):
         self.input = a
 
 
+class MockModel_(BaseModel):
+    __REQUIRED_DATA__ = ["x"]
+    __REQUIRED_LABELS__ = ["y"]
+
+    def __init__(self):
+        super(MockModel_, self).__init__(DictConfig({"conv_type": "Dummy"}))
+
+        self._channels = [12, 12, 12, 17]
+        self.nn = MLP(self._channels)
+
+    def set_input(self, a):
+        self.input = a
+
+
 class TestBaseModel(unittest.TestCase):
     def test_getinput(self):
         model = MockModel()
@@ -71,6 +86,21 @@ def test_enable_dropout_eval(self):
             self.assertEqual(model.nn[i][1].training, True)
             self.assertEqual(model.nn[i][2].training, False)
 
+    def test_load_pretrained_model(self):
+        """
+        test load_state_dict_with_same_shape
+        """
+        model1 = MockModel()
+        model2 = MockModel_()
+
+        w1 = model1.state_dict()
+
+        model2.load_state_dict_with_same_shape(w1)
+        w2 = model2.state_dict()
+        for k, p in w2.items():
+            if "nn.2." not in k:
+                torch.testing.assert_allclose(w1[k], p)
+
     def test_accumulated_gradient(self):
         params = load_model_config("segmentation", "pointnet2", "pointnet2ms")
         config_training = OmegaConf.load(os.path.join(DIR, "test_config/training_config.yaml"))

torch_points3d/models/base_model.py

@@ -1,6 +1,7 @@
 from collections import OrderedDict
 from abc import abstractmethod
 from typing import Optional, Dict, Any, List
+import os
 import torch
 from torch.optim.optimizer import Optimizer
 from torch.optim.lr_scheduler import _LRScheduler
@@ -133,6 +134,24 @@ def set_input(self, input, device):
         """
         raise NotImplementedError
 
+    def load_state_dict_with_same_shape(self, weights, strict=False):
+        model_state = self.state_dict()
+        filtered_weights = {k: v for k, v in weights.items() if k in model_state and v.size() == model_state[k].size()}
+        log.info("Loading weights:" + ", ".join(filtered_weights.keys()))
+        self.load_state_dict(filtered_weights, strict=strict)
+
+    def set_pretrained_weights(self):
+        path_pretrained = getattr(self.opt, "path_pretrained", None)
+        weight_name = getattr(self.opt, "weight_name", "latest")
+
+        if path_pretrained is not None:
+            if not os.path.exists(path_pretrained):
+                log.warning("The path does not exist, it will not load any model")
+            else:
+                log.info("load pretrained weights from {}".format(path_pretrained))
+                m = torch.load(path_pretrained)["models"][weight_name]
+                self.load_state_dict_with_same_shape(m, strict=False)
+
     def get_labels(self):
         """ returns a trensor of size ``[N_points]`` where each value is the label of a point
         """

torch_points3d/trainer.py

@@ -93,9 +93,11 @@ def _initialize_trainer(self):
                 )
             self._model: BaseModel = instantiate_model(copy.deepcopy(self._cfg), self._dataset)
             self._model.instantiate_optimizers(self._cfg)
+            self._model.set_pretrained_weights()
         self._checkpoint.dataset_properties = self._dataset.used_properties
 
         log.info(self._model)
+
         self._model.log_optimizers()
         log.info("Model size = %i", sum(param.numel() for param in self._model.parameters() if param.requires_grad))