🔄 CFlow: Switch soft permutation to false by default to speed up training.

anomalib/models/cflow/backbone.py

@@ -73,14 +73,19 @@ def subnet_fc(dims_in: int, dims_out: int):
     return nn.Sequential(nn.Linear(dims_in, 2 * dims_in), nn.ReLU(), nn.Linear(2 * dims_in, dims_out))
 
 
-def cflow_head(condition_vector: int, coupling_blocks: int, clamp_alpha: float, n_features: int) -> SequenceINN:
+def cflow_head(
+    condition_vector: int, coupling_blocks: int, clamp_alpha: float, n_features: int, permute_soft: bool = False
+) -> SequenceINN:
     """Create invertible decoder network.
 
     Args:
         condition_vector (int): length of the condition vector
         coupling_blocks (int): number of coupling blocks to build the decoder
         clamp_alpha (float): clamping value to avoid exploding values
         n_features (int): number of decoder features
+        permute_soft (bool): Whether to sample the permutation matrix :math:`R` from :math:`SO(N)`,
+            or to use hard permutations instead. Note, ``permute_soft=True`` is very slow
+            when working with >512 dimensions.
 
     Returns:
         SequenceINN: decoder network block
@@ -95,6 +100,6 @@ def cflow_head(condition_vector: int, coupling_blocks: int, clamp_alpha: float,
             subnet_constructor=subnet_fc,
             affine_clamping=clamp_alpha,
             global_affine_type="SOFTPLUS",
-            permute_soft=True,
+            permute_soft=permute_soft,
         )
     return coder

anomalib/models/cflow/config.yaml

@@ -24,6 +24,7 @@ model:
   condition_vector: 128
   coupling_blocks: 8
   clamp_alpha: 1.9
+  soft_permutation: false
   lr: 0.0001
   early_stopping:
     patience: 2

anomalib/models/cflow/model.py

@@ -141,7 +141,13 @@ def __init__(self, hparams: Union[DictConfig, ListConfig]):
         self.pool_dims = self.encoder.out_dims
         self.decoders = nn.ModuleList(
             [
-                cflow_head(self.condition_vector, hparams.model.coupling_blocks, hparams.model.clamp_alpha, pool_dim)
+                cflow_head(
+                    condition_vector=self.condition_vector,
+                    coupling_blocks=hparams.model.coupling_blocks,
+                    clamp_alpha=hparams.model.clamp_alpha,
+                    n_features=pool_dim,
+                    permute_soft=hparams.model.soft_permutation,
+                )
                 for pool_dim in self.pool_dims
             ]
         )