feat: add possibility to return no kl, save it as attribute

Author: piEsposito

bayesian_torch/layers/variational_layers/conv_variational.py

@@ -112,6 +112,8 @@ def __init__(self,
         self.posterior_rho_init = posterior_rho_init,
         self.bias = bias
 
+        self.kl = 0
+
         self.mu_kernel = Parameter(
             torch.Tensor(out_channels, in_channels // groups, kernel_size))
         self.rho_kernel = Parameter(
@@ -160,7 +162,7 @@ def init_parameters(self):
             self.rho_bias.data.normal_(mean=self.posterior_rho_init[0],
                                        std=0.1)
 
-    def forward(self, input):
+    def forward(self, input, return_kl=True):
         sigma_weight = torch.log1p(torch.exp(self.rho_kernel))
         eps_kernel = self.eps_kernel.data.normal_()
         weight = self.mu_kernel + (sigma_weight * eps_kernel)
@@ -182,7 +184,11 @@ def forward(self, input):
         else:
             kl = kl_weight
 
-        return out, kl
+        self.kl = kl
+
+        if return_kl:
+            return out, kl
+        return out
 
 
 class Conv2dReparameterization(BaseVariationalLayer_):
@@ -239,6 +245,8 @@ def __init__(self,
         self.posterior_rho_init = posterior_rho_init,
         self.bias = bias
 
+        self.kl = 0
+
         self.mu_kernel = Parameter(
             torch.Tensor(out_channels, in_channels // groups, kernel_size,
                          kernel_size))
@@ -292,7 +300,7 @@ def init_parameters(self):
             self.rho_bias.data.normal_(mean=self.posterior_rho_init[0],
                                        std=0.1)
 
-    def forward(self, input):
+    def forward(self, input, return_kl=True):
         sigma_weight = torch.log1p(torch.exp(self.rho_kernel))
         eps_kernel = self.eps_kernel.data.normal_()
         weight = self.mu_kernel + (sigma_weight * eps_kernel)
@@ -313,8 +321,12 @@ def forward(self, input):
             kl = kl_weight + kl_bias
         else:
             kl = kl_weight
+        
+        self.kl = kl
 
-        return out, kl
+        if return_kl:
+            return out, kl
+        return out
 
 
 class Conv3dReparameterization(BaseVariationalLayer_):
@@ -371,6 +383,8 @@ def __init__(self,
         self.posterior_rho_init = posterior_rho_init,
         self.bias = bias
 
+        self.kl = 0
+
         self.mu_kernel = Parameter(
             torch.Tensor(out_channels, in_channels // groups, kernel_size,
                          kernel_size, kernel_size))
@@ -424,7 +438,7 @@ def init_parameters(self):
             self.rho_bias.data.normal_(mean=self.posterior_rho_init[0],
                                        std=0.1)
 
-    def forward(self, input):
+    def forward(self, input, return_kl=True):
         sigma_weight = torch.log1p(torch.exp(self.rho_kernel))
         eps_kernel = self.eps_kernel.data.normal_()
         weight = self.mu_kernel + (sigma_weight * eps_kernel)
@@ -446,7 +460,11 @@ def forward(self, input):
         else:
             kl = kl_weight
 
-        return out, kl
+        self.kl = kl
+
+        if return_kl:
+            return out, kl
+        return out
 
 
 class ConvTranspose1dReparameterization(BaseVariationalLayer_):
@@ -504,6 +522,8 @@ def __init__(self,
         self.posterior_rho_init = posterior_rho_init,
         self.bias = bias
 
+        self.kl = 0
+
         self.mu_kernel = Parameter(
             torch.Tensor(in_channels, out_channels // groups, kernel_size))
         self.rho_kernel = Parameter(
@@ -552,7 +572,7 @@ def init_parameters(self):
             self.rho_bias.data.normal_(mean=self.posterior_rho_init[0],
                                        std=0.1)
 
-    def forward(self, input):
+    def forward(self, input, return_kl=True):
         sigma_weight = torch.log1p(torch.exp(self.rho_kernel))
         eps_kernel = self.eps_kernel.data.normal_()
         weight = self.mu_kernel + (sigma_weight * eps_kernel)
@@ -575,7 +595,11 @@ def forward(self, input):
         else:
             kl = kl_weight
 
-        return out, kl
+        self.kl = kl
+
+        if return_kl:
+            return out, kl
+        return out
 
 
 class ConvTranspose2dReparameterization(BaseVariationalLayer_):
@@ -633,6 +657,8 @@ def __init__(self,
         self.posterior_rho_init = posterior_rho_init,
         self.bias = bias
 
+        self.kl = 0
+
         self.mu_kernel = Parameter(
             torch.Tensor(in_channels, out_channels // groups, kernel_size,
                          kernel_size))
@@ -686,7 +712,7 @@ def init_parameters(self):
             self.rho_bias.data.normal_(mean=self.posterior_rho_init[0],
                                        std=0.1)
 
-    def forward(self, input):
+    def forward(self, input, return_kl=True):
         sigma_weight = torch.log1p(torch.exp(self.rho_kernel))
         eps_kernel = self.eps_kernel.data.normal_()
         weight = self.mu_kernel + (sigma_weight * eps_kernel)
@@ -709,7 +735,11 @@ def forward(self, input):
         else:
             kl = kl_weight
 
-        return out, kl
+        self.kl = kl
+
+        if return_kl:
+            return out, kl
+        return out
 
 
 class ConvTranspose3dReparameterization(BaseVariationalLayer_):
@@ -768,6 +798,8 @@ def __init__(self,
         self.posterior_rho_init = posterior_rho_init,
         self.bias = bias
 
+        self.kl = 0
+
         self.mu_kernel = Parameter(
             torch.Tensor(in_channels, out_channels // groups, kernel_size,
                          kernel_size, kernel_size))
@@ -821,7 +853,7 @@ def init_parameters(self):
             self.rho_bias.data.normal_(mean=self.posterior_rho_init[0],
                                        std=0.1)
 
-    def forward(self, input):
+    def forward(self, input, return_kl=True):
         sigma_weight = torch.log1p(torch.exp(self.rho_kernel))
         eps_kernel = self.eps_kernel.data.normal_()
         weight = self.mu_kernel + (sigma_weight * eps_kernel)
@@ -844,4 +876,8 @@ def forward(self, input):
         else:
             kl = kl_weight
 
-        return out, kl
+        self.kl = kl
+
+        if return_kl:
+            return out, kl
+        return out

bayesian_torch/layers/variational_layers/linear_variational.py

@@ -83,6 +83,8 @@ def __init__(self,
         self.posterior_rho_init = posterior_rho_init,
         self.bias = bias
 
+        self.kl = 0
+
         self.mu_weight = Parameter(torch.Tensor(out_features, in_features))
         self.rho_weight = Parameter(torch.Tensor(out_features, in_features))
         self.register_buffer('eps_weight',
@@ -124,7 +126,7 @@ def init_parameters(self):
             self.rho_bias.data.normal_(mean=self.posterior_rho_init[0],
                                        std=0.1)
 
-    def forward(self, input):
+    def forward(self, input, return_kl=True):
         sigma_weight = torch.log1p(torch.exp(self.rho_weight))
         weight = self.mu_weight + \
             (sigma_weight * self.eps_weight.data.normal_())
@@ -143,5 +145,9 @@ def forward(self, input):
             kl = kl_weight + kl_bias
         else:
             kl = kl_weight
+            
+        self.kl = kl
 
-        return out, kl
+        if return_kl:
+            return out, kl
+        return out

bayesian_torch/layers/variational_layers/rnn_variational.py

@@ -77,6 +77,8 @@ def __init__(self,
         self.posterior_rho_init = posterior_rho_init,
         self.bias = bias
 
+        self.kl = kl
+
         self.ih = LinearReparameterization(
             prior_mean=prior_mean,
             prior_variance=prior_variance,
@@ -95,7 +97,7 @@ def __init__(self,
             out_features=out_features * 4,
             bias=bias)
 
-    def forward(self, X, hidden_states=None):
+    def forward(self, X, hidden_states=None, return_kl=True):
 
         batch_size, seq_size, _ = X.size()
 
@@ -140,4 +142,8 @@ def forward(self, X, hidden_states=None):
         hidden_seq = hidden_seq.transpose(0, 1).contiguous()
         c_ts = c_ts.transpose(0, 1).contiguous()
 
-        return hidden_seq, (hidden_seq, c_ts), kl
+        self.kl = kl
+
+        if return_kl:
+            return hidden_seq, (hidden_seq, c_ts), kl
+        return hidden_seq, (hidden_seq, c_ts)