feat: parameter groups & splatting monads

+ dsl updates (tanh, sqrt, roll)
! fix: bug in dsl's sub order when using scalar numbers
+ new torch loader kwargs
+ image torchmetrics
+ quat monads
+ adopt optimizer
+ updated parameter selection to support different optimizer kwargs for different params
! fix: zip iterator
+ add rerun config logging
! fix: scheduling for fitting pipelines
+ mode for rerun scalars logging

Author: nmvrs

hydra_plugins/moai_dsl_plugin/moai_dsl_plugin.py

@@ -92,8 +92,12 @@
         | "asin" "(" expr ")"               -> asin
         | "tan" "(" name ")"                -> tan
         | "tan" "(" expr ")"                -> tan
+        | "tanh" "(" name ")"               -> tanh
+        | "tanh" "(" expr ")"               -> tanh
         | "atan" "(" name ")"               -> atan
         | "atan" "(" expr ")"               -> atan
+        | "sqrt" "(" name ")"               -> sqrt
+        | "sqrt" "(" expr ")"               -> sqrt
         | "deg" "(" name ")"                -> rad2deg
         | "deg" "(" expr ")"                -> rad2deg
         | "rad" "(" name ")"                -> deg2rad
@@ -108,6 +112,7 @@
         | "transpose" "(" name "," SIGNED_INT ("," SIGNED_INT)* ")"  -> transpose
         | "flatten" "(" name "," SIGNED_INT ["," SIGNED_INT] ")"  -> flatten
         | "repeat_interleave" "(" name "," SIGNED_INT "," SIGNED_INT ")" -> repeat
+        | "roll" "(" name "," SIGNED_INT "," SIGNED_INT ")" -> roll
         | "zeros" "(" name ")"              -> zeros_like
         | "ones" "(" name ")"               -> ones_like
         | "rand" "(" name ")"               -> rand_like

moai/conf/data/test/loader/torch.yaml

@@ -5,4 +5,8 @@ batch_size: ${oc.select:experiment.batch_size,1}
 shuffle: false
 num_workers: ${oc.select:experiment.workers,0}
 pin_memory: false
-drop_last: false
+drop_last: false
+prefetch_factor: 0
+persistent_workers: false
+pin_memory_device: ""
+# in_order: true

moai/conf/data/train/loader/torch.yaml

@@ -5,4 +5,8 @@ batch_size: ${oc.select:experiment.batch_size,1}
 shuffle: true
 num_workers: ${oc.select:experiment.workers,0}
 pin_memory: false
-drop_last: false
+drop_last: false
+prefetch_factor: 0
+persistent_workers: false
+pin_memory_device: ""
+# in_order: true

moai/conf/model/metrics/torchmetrics/image/lpips.yaml

@@ -0,0 +1,7 @@
+# @package model.metrics.learned_perceptual_image_patch_similarity
+
+_target_: moai.validation.torchmetric.LPIPS
+type: LearnedPerceptualImagePatchSimilarity
+net_type: alex # vgg, squeeze
+normalize: false
+reduction: mean # sum, none

moai/conf/model/metrics/torchmetrics/image/psnr.yaml

@@ -0,0 +1,8 @@
+# @package model.metrics.peak_signal_noise_ratio
+
+_target_: moai.validation.torchmetric.ImageTorchMetric
+type: PeakSignalNoiseRatio
+data_range: 1.0
+base: 10.0
+reduction: elementwise_mean # sum, none
+dim: null

moai/conf/model/metrics/torchmetrics/image/ssim.yaml

@@ -0,0 +1,11 @@
+# @package model.metrics.structural_similarity_index_measure
+
+_target_: moai.validation.torchmetric.ImageTorchMetric
+type: StructuralSimilarityIndexMeasure
+gaussian_kernel: true
+sigma: 1.5
+kernel_size: 11
+reduction: elementwise_mean # sum, none
+data_range: 1.0
+k1: 0.01
+k2: 0.03

moai/conf/model/monads/generation/splats/gaussian_splat_parameters.yaml

@@ -0,0 +1,9 @@
+# @package model.monads.gaussian_splat_parameters
+
+_target_: moai.monads.generation.splats.gaussian_parameters.GaussianSplattingParameters
+num_splats: ???
+num_sets: 1
+skip_positions: false
+scale: 1.0
+opacity: null
+sh_degree: 0

moai/conf/model/monads/generation/tensor/clone.yaml

@@ -1,3 +1,3 @@
-# @package model.monads._name_
+# @package model.monads.clone
 
 _target_: moai.monads.generation.tensor.Clone

moai/conf/model/monads/generation/tensor/named_params.yaml

@@ -1,4 +1,4 @@
-# @package model.monads._name_
+# @package model.monads.named_params
 
 _target_: moai.monads.generation.tensor.Parameters
 parameters: ???

moai/conf/model/monads/generation/tensor/npy.yaml

@@ -1,3 +1,3 @@
-# @package model.monads._name_
+# @package model.monads.npy
 
 _target_: moai.monads.generation.tensor.Npy

moai/conf/model/monads/generation/tensor/ones.yaml

@@ -1,4 +1,4 @@
-# @package model.monads._name_
+# @package model.monads.ones
 
 _target_: moai.monads.generation.tensor.Ones
 shape: ???

moai/conf/model/monads/geometry/rotation/quaternion_composition.yaml

@@ -0,0 +1,4 @@
+# @package model.monads.quaternion_composition
+
+_target_: moai.monads.geometry.rotations.quaternion.QuaternionComposition
+normalize: false

moai/conf/model/monads/geometry/rotation/unitquat_to_rotmat.yaml

@@ -1,3 +1,3 @@
-# @package model.monads._name_
+# @package model.monads.unitquat_to_rotmat
 
 _target_: moai.monads.geometry.rotations.Quaternion2RotationMatrix

moai/conf/model/monads/math/clamp.yaml

@@ -1,4 +1,4 @@
-# @package model.monads._name_
+# @package model.monads.clamp
 
 _target_: moai.monads.math.Clamp
 min_value: 0.0

moai/conf/model/monads/render/splats/gaussian_splat_rasterizer.yaml

@@ -0,0 +1,9 @@
+# @package model.monads.gaussian_splat_rasterizer
+
+_target_: moai.monads.render.splats.rasterize.GaussianSplatRasterizer
+prefiltered: false
+debug: false
+scale_modifier: 1.0
+width: null
+height: null
+background_color: black

moai/conf/model/monads/render/splats/gaussian_splat_rasterizer_ashawkey.yaml

@@ -0,0 +1,9 @@
+# @package model.monads.gaussian_splat_rasterizer
+
+_target_: moai.monads.render.splats.rasterize_ashawkey.GaussianSplatRasterizer
+prefiltered: false
+debug: false
+scale_modifier: 1.0
+width: null
+height: null
+background_color: black

moai/conf/model/monads/sampling/remap.yaml

@@ -0,0 +1,6 @@
+# @package model.monads.remap
+
+_target_: moai.monads.sampling.torch.remap.Remap
+align_corners: null
+sample_mode: bilinear
+padding_mode: zeros

moai/conf/model/monads/tensor/detach.yaml

@@ -1,3 +1,3 @@
-# @package model.monads._name_
+# @package model.monads.detach
 
 _target_: moai.monads.tensor.Detach

moai/conf/model/monitors/visualization/rerun/config.yaml

@@ -0,0 +1,4 @@
+# @package model.monitors.config
+
+_target_: moai.visualization.rerun.config.Config
+path: '/config'

moai/conf/model/objectives/image/structural_dissimilarity.yaml

@@ -0,0 +1,5 @@
+# @package model.objectives.structural_dissimilarity
+
+_target_: moai.supervision.objectives.image.ssim.StructuralDissimilarity
+window_size: 7
+dynamic_range: 1.0

moai/conf/model/objectives/mesh/surface_feature_consistency.yaml

@@ -0,0 +1,4 @@
+# @package model.objectives.surface_feature_consistency
+
+_target_: moai.supervision.objectives.mesh.consistency.SurfaceFeatureConsistency
+order: 2

moai/conf/model/objectives/passthrough.yaml

@@ -1,4 +1,4 @@
 # @package model.objectives.passthrough
 
-_target_: moai.supervision.objectives.Passthrough
+_target_: moai.supervision.objectives.passthrough.Passthrough
 mode: minimize # one of [minimize, maximize]

moai/conf/model/objectives/splats/long_and_thin_penalty.yaml

@@ -0,0 +1,5 @@
+# @package model.objectives.long_and_thin_penalty
+
+_target_: moai.supervision.objectives.splats.scaling.LongAndThinPenalty
+max_scale: 0.008
+scale_ratio: 10.0

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

@@ -0,0 +1,7 @@
+# @package model.parameters.groups.model3
+
+_target_: moai.parameters.selectors.model.ModelParameterSelector
+modules: null # optional
+monads: null # optional
+parameters: null # optional
+force_grad: true

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

@@ -0,0 +1,7 @@
+# @package model.parameters.groups.model4
+
+_target_: moai.parameters.selectors.model.ModelParameterSelector
+modules: null # optional
+monads: null # optional
+parameters: null # optional
+force_grad: true

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

@@ -0,0 +1,7 @@
+# @package model.parameters.groups.model5
+
+_target_: moai.parameters.selectors.model.ModelParameterSelector
+modules: null # optional
+monads: null # optional
+parameters: null # optional
+force_grad: true

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

@@ -0,0 +1,12 @@
+# @package model.parameters.groups.model_groups
+
+_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_groups1.yaml

@@ -0,0 +1,12 @@
+# @package model.parameters.groups.model_groups1
+
+_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_groups2.yaml

@@ -0,0 +1,12 @@
+# @package model.parameters.groups.model_groups2
+
+_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_groups3.yaml

@@ -0,0 +1,12 @@
+# @package model.parameters.groups.model_groups3
+
+_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/initializers/default.yaml

@@ -1,3 +1,3 @@
 # @package model.parameters.initializers.default
 
-_target_: moai.parameters.initialization.default.Default
+_target_: moai.parameters.initialization.default.Default

moai/conf/model/parameters/optimizers/adopt.yaml

@@ -0,0 +1,13 @@
+# @package model.parameters.optimizers.adopt
+
+_target_: moai.parameters.optimization.optimizers.adopt.ADOPT
+lr: 0.0001
+betas: [0.9, 0.999]
+eps: 1e-6 # 1e-8
+weight_decay: 0.0
+decouple: false
+foreach: null
+maximize: false
+capturable: false
+differentiable: false
+fused: null

moai/core/execution/expression.py

@@ -1,4 +1,5 @@
 import dataclasses
+import numbers
 import typing
 
 import benedict
@@ -74,21 +75,27 @@ def forward(self, td, tmp) -> None:
 @dataclasses.dataclass(repr=False, unsafe_hash=True)
 class BinaryOperationScalar(torch.nn.Module):
     operation: str
-    lhs: str
-    rhs: typing.Union[float, int]
+    lhs: typing.Union[float, int, str]
+    rhs: typing.Union[float, int, str]
     index: int
 
     def __post_init__(self):
         super().__init__()
         self.op = getattr(torch, self.operation)
+        self.is_lhs_scalar = isinstance(self.lhs, numbers.Number)
 
     def __repr__(self):
         return f"{self.operation}:{self.lhs},{self.rhs}"
 
     def forward(self, td, tmp) -> None:
-        tmp[f"result{self.index}"] = self.op(
-            toolz.get_in(self.lhs.split("."), tmp), self.rhs
-        )
+        if self.is_lhs_scalar:
+            tmp[f"result{self.index}"] = self.op(
+                -toolz.get_in(self.rhs.split("."), tmp), -self.lhs
+            )
+        else:
+            tmp[f"result{self.index}"] = self.op(
+                toolz.get_in(self.lhs.split("."), tmp), self.rhs
+            )
 
 
 @dataclasses.dataclass(repr=False, unsafe_hash=True)
@@ -283,7 +290,7 @@ def _binary(self, name, lhs, rhs):
         if rhs is None:
             rhs = self.results.pop()
         if not isinstance(lhs, str):
-            m = BinaryOperationScalar(name, rhs, lhs, self.index)
+            m = BinaryOperationScalar(name, lhs, rhs, self.index)
         elif not isinstance(rhs, str):
             m = BinaryOperationScalar(name, lhs, rhs, self.index)
         else:
@@ -572,6 +579,10 @@ def acos(self, key):
     def tan(self, key):
         pass
 
+    @unary
+    def tanh(self, key):
+        pass
+
     @unary
     def atan(self, key):
         pass
@@ -581,7 +592,7 @@ def abs(self, key):
         pass
 
     @unary
-    def abs(self, key):
+    def sqrt(self, key):
         pass
 
     @unary
@@ -716,6 +727,10 @@ def repeat(self, key, *dims):
         dims = list(map(int, dims))
         self._transform_operation("repeat_interleave", key, dims)
 
+    def roll(self, key, shift, dim):
+        key = self.extract(key)  # NOTE: only supports single shift/dim
+        self._transform_operation("roll", key, [int(shift), int(dim)])
+
     def unsqueeze(self, key, *dims):
         if not isinstance(key, str) or isinstance(key, Token):  # NOTE: is lark.Tree
             key = self.extract(key)

moai/core/execution/metrics.py

@@ -62,7 +62,7 @@ def __init__(
             extra_params = set(metric_kwargs.keys()) - set(sig_params) - set([C._OUT_])
             if extra_params:
                 log.error(
-                    f"The parameters [{extra_params}] are not part of the `{key}` metric signature."
+                    f"The parameters [{extra_params}] are not part of the `{key}` metric signature ({list(sig.parameters.keys())})."
                 )
             metric_kwargs = mic._dict_of_lists_to_list_of_dicts(metric_kwargs)
             for j, params in enumerate(metric_kwargs):

moai/core/execution/models.py

@@ -43,7 +43,7 @@ def __init__(
             extra_params = set(model_params.keys()) - set(sig_params) - set([C._OUT_])
             if extra_params:
                 log.error(
-                    f"The parameters [{extra_params}] are not part of the `{key}` model signature."
+                    f"The parameters [{extra_params}] are not part of the `{key}` model signature ({list(sig.parameters.keys())})."
                 )
             model_params = mic._dict_of_lists_to_list_of_dicts(model_params)
             for j, params in enumerate(model_params):

moai/core/execution/monads.py

@@ -69,7 +69,7 @@ def __init__(
             extra_params = set(graph_params.keys()) - set(sig_params) - set([C._OUT_])
             if extra_params:
                 log.error(
-                    f"The parameters [{extra_params}] are not part of the `{key}` monad signature."
+                    f"The parameters [{extra_params}] are not part of the `{key}` monad signature ({list(sig.parameters.keys())})."
                 )
             graph_params = mic._dict_of_lists_to_list_of_dicts(graph_params)
             for j, params in enumerate(graph_params):