Upload nets.py

nets.py

@@ -14,6 +14,8 @@
 
 """Deep networks."""
 
+from copy import deepcopy
+
 import numpy as np
 import torch
 import torch.nn.functional as F
@@ -131,10 +133,14 @@ class EnsembleFC(nn.Module):
         return torch.add(wx, self.bias[:, None, None, :])  # w times x + b
 
 
-class EnsembleModel(nn.Module):
+def get_params(model):
+    return torch.cat([p.view(-1) for p in model.parameters()])
+
+
+class _EnsembleModel(nn.Module):
     def __init__(self, encoding_dim, num_ensemble, hidden_dim=128, activation="relu", dtype=torch.float32) -> None:
         # super().__init__(encoding_dim, hidden_dim, activation)
-        super(EnsembleModel, self).__init__()
+        super(_EnsembleModel, self).__init__()
         self.num_ensemble = num_ensemble
         self.hidden_dim = hidden_dim
         self.output_dim = 1
@@ -152,23 +158,34 @@ class EnsembleModel(nn.Module):
         else:
             raise ValueError(f"Unknown activation {activation}")
 
-    def get_params(self) -> torch.Tensor:
-        params = []
-        for pp in list(self.parameters()):
-            params.append(pp.view(-1))
-        return torch.cat(params)
-
     def forward(self, encoding: torch.Tensor) -> torch.Tensor:
         x = self.activation(self.nn1(encoding))
         x = self.activation(self.nn2(x))
         score = self.nn_out(x)
         return score
 
-    def init(self):
-        self.init_params = self.get_params().data.clone()
-        if torch.cuda.is_available():
-            self.init_params = self.init_params.cuda()
-
     def regularization(self):
         """Prior towards independent initialization."""
         return ((self.get_params() - self.init_params) ** 2).mean()
+
+
+class EnsembleModel(nn.Module):
+    def __init__(self, encoding_dim, num_ensemble, hidden_dim=128, activation="relu", dtype=torch.float32) -> None:
+        super(EnsembleModel, self).__init__()
+        self.encoding_dim = encoding_dim
+        self.num_ensemble = num_ensemble
+        self.hidden_dim = hidden_dim
+        self.model = _EnsembleModel(encoding_dim, num_ensemble, hidden_dim, activation, dtype)
+        self.reg_model = deepcopy(self.model)  # only used for regularization
+        # freeze the reg model
+        for param in self.reg_model.parameters():
+            param.requires_grad = False
+
+    def forward(self, encoding: torch.Tensor) -> torch.Tensor:
+        return self.model(encoding)
+
+    def regularization(self):
+        """Prior towards independent initialization."""
+        model_params = get_params(self.model)
+        reg_params = get_params(self.reg_model).detach()
+        return ((model_params - reg_params) ** 2).mean()