Coupling flow working but not performing well with current clamping

bayesflow-org · May 16, 2024 · 51cb6a7 · 51cb6a7
1 parent 3edb1e6
commit 51cb6a7
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/bayesflow/experimental/networks/coupling_flow/transforms/affine_transform.py b/bayesflow/experimental/networks/coupling_flow/transforms/affine_transform.py
@@ -9,7 +9,7 @@
 
 class AffineTransform(Transform):
 
-    def __init__(self, clamp_factor=10.0, **kwargs):
+    def __init__(self, clamp_factor=5.0, **kwargs):
         super().__init__(**kwargs)
         self.clamp_factor = clamp_factor
 
@@ -19,8 +19,10 @@ def split_parameters(self, parameters: Tensor) -> dict[str, Tensor]:
         return {"scale": scale, "shift": shift}
 
     def constrain_parameters(self, parameters: dict[str, Tensor]) -> dict[str, Tensor]:
-        shift = math.log(math.e - 1)
-        parameters["scale"] = self.clamp_factor * ops.sigmoid(ops.softplus(parameters["scale"] + shift))
+        # shift = math.log(math.e - 1)
+        s = parameters["scale"]
+        # parameters["scale"] = self.clamp_factor * ops.sigmoid(ops.softplus(parameters["scale"] + shift))
+        parameters["scale"] = 1 / (1 + ops.exp(-s)) * ops.sqrt(1 + ops.abs(s + self.clamp_factor))
 
         return parameters
 

diff --git a/bayesflow/experimental/networks/resnet/hidden_block.py b/bayesflow/experimental/networks/resnet/hidden_block.py
@@ -6,7 +6,7 @@ class ConfigurableHiddenBlock(keras.layers.Layer):
     def __init__(
         self,
         num_units,
-        activation="gelu",
+        activation="relu",
         residual=True,
         dropout_rate=0.05,
         spectral_norm=False,