🔨 Rework stochastic models

ENSTA-U2IS-AI · Jul 1, 2023 · b4a42b7 · b4a42b7
1 parent 84449a3
commit b4a42b7
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/torch_uncertainty/models/lenet.py b/torch_uncertainty/models/lenet.py
@@ -6,7 +6,7 @@
 
 from ..layers.bayesian_layers import BayesConv2d, BayesLinear
 from ..layers.packed_layers import PackedConv2d, PackedLinear
-from .variational_model import Variational
+from .utils import Stochastic
 
 
 class LeNet(nn.Module):
@@ -41,7 +41,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return out
 
 
-@Variational
+@Stochastic
 class BayesianLeNet(LeNet):
     pass
 

diff --git a/...h_uncertainty/models/variational_model.py → torch_uncertainty/models/utils.py b/...h_uncertainty/models/variational_model.py → torch_uncertainty/models/utils.py
@@ -1,7 +1,13 @@
 from ..layers.bayesian_layers import bayesian_modules
 
 
-def Variational(Model):
+def Stochastic(Model):
+    """Decorator for stochastic models. When applied to a model, it adds the
+    freeze and unfreeze methods to the model. Use freeze to freeze the
+    stochastic layers and obtain deterministic outputs. Use unfreeze to
+    unfreeze the stochastic layers and obtain stochastic outputs.
+    """
+
     def freeze(self):
         for module in self.modules():
             if isinstance(module, bayesian_modules):