feat: add extra model groups

Author: nmvrs

moai/conf/model/parameters/groups/model_groups4.yaml

@@ -0,0 +1,12 @@
+# @package model.parameters.groups.model_groups4
+
+_target_: moai.parameters.selectors.model_groups.ModelGroupParameterSelector
+groups: ???
+# each group (key-dict pair) contains
+  # modules: null # optional
+  # monads: null # optional
+  # parameters: null # optional
+  # force_grad: true
+# and extra optimizer relaed params
+  # lr: float
+  # etc.

moai/conf/model/parameters/groups/model_groups5.yaml

@@ -0,0 +1,12 @@
+# @package model.parameters.groups.model_groups5
+
+_target_: moai.parameters.selectors.model_groups.ModelGroupParameterSelector
+groups: ???
+# each group (key-dict pair) contains
+  # modules: null # optional
+  # monads: null # optional
+  # parameters: null # optional
+  # force_grad: true
+# and extra optimizer relaed params
+  # lr: float
+  # etc.

moai/core/model.py

@@ -289,6 +289,7 @@ def closure(tensors, index, steps, stage, optimizer, objective):
                     optimizer,
                 )
             call._call_strategy_hook(self.trainer, "backward", loss, optimizer)
+            # NOTE: https://pytorch.org/docs/main/generated/torch.optim.Optimizer.zero_grad.html#torch.optim.Optimizer.zero_grad
             self.optimization_step += 1
             if monitor := toolz.get_in(
                 [C._FIT_, C._OPTIMIZATION_STEP_, stage], self.monitor

moai/monads/render/nvdiffrast/render.py

@@ -64,6 +64,7 @@ def forward(
         attributes: typing.Optional[torch.Tensor] = None,
     ) -> typing.Dict[str, torch.Tensor]:
         nz_indices = torch.nonzero(triangle_ids[..., -1], as_tuple=True)
+        # NOTE: reconsider nonzero as it seems to be a blocking call
         out = {
             "indices": dict((str(i), t) for i, t in enumerate(nz_indices)),
             "triangles": {"id": triangle_ids[nz_indices]},

moai/nn/convolution/torch/standard.py

@@ -27,6 +27,7 @@ def __init__(
             out_channels=out_features,
             **convolution_params  # TODO: either merge kwargs here or in the factory method
         )
+        # https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html#disable-bias-for-convolutions-directly-followed-by-a-batch-norm
         self.activation = mia.make_activation(
             features=out_features, activation_type=activation_type, **activation_params
         )