Merge pull request #9 from piEsposito/main

Let the models return prediction only, saving KL Divergence as an attribute

Author: ranganathkrishnan

bayesian_torch/layers/flipout_layers/conv_flipout.py

@@ -100,6 +100,8 @@ def __init__(self,
         self.posterior_rho_init = posterior_rho_init
         self.bias = bias
 
+        self.kl = 0
+
         self.mu_kernel = nn.Parameter(
             torch.Tensor(out_channels, in_channels // groups, kernel_size))
         self.rho_kernel = nn.Parameter(
@@ -150,7 +152,7 @@ def init_parameters(self):
             self.prior_bias_mu.data.fill_(self.prior_mean)
             self.prior_bias_sigma.data.fill_(self.prior_variance)
 
-    def forward(self, x):
+    def forward(self, x, return_kl=True):
 
         # linear outputs
         outputs = F.conv1d(x,
@@ -191,8 +193,11 @@ def forward(self, x):
                                      dilation=self.dilation,
                                      groups=self.groups) * sign_output
 
+        self.kl = kl
         # returning outputs + perturbations
-        return outputs + perturbed_outputs, kl
+        if return_kl:
+            return outputs + perturbed_outputs, kl
+        return outputs + perturbed_outputs
 
 
 class Conv2dFlipout(BaseVariationalLayer_):
@@ -244,6 +249,8 @@ def __init__(self,
         self.posterior_rho_init = posterior_rho_init
         self.bias = bias
 
+        self.kl = 0
+
         self.mu_kernel = nn.Parameter(
             torch.Tensor(out_channels, in_channels // groups, kernel_size,
                          kernel_size))
@@ -299,7 +306,7 @@ def init_parameters(self):
             self.prior_bias_mu.data.fill_(self.prior_mean)
             self.prior_bias_sigma.data.fill_(self.prior_variance)
 
-    def forward(self, x):
+    def forward(self, x, return_kl=True):
 
         # linear outputs
         outputs = F.conv2d(x,
@@ -340,8 +347,11 @@ def forward(self, x):
                                      dilation=self.dilation,
                                      groups=self.groups) * sign_output
 
+        self.kl = kl
         # returning outputs + perturbations
-        return outputs + perturbed_outputs, kl
+        if return_kl:
+            return outputs + perturbed_outputs, kl
+        return outputs + perturbed_outputs
 
 
 class Conv3dFlipout(BaseVariationalLayer_):
@@ -388,6 +398,8 @@ def __init__(self,
         self.groups = groups
         self.bias = bias
 
+        self.kl = 0
+
         self.prior_mean = prior_mean
         self.prior_variance = prior_variance
         self.posterior_mu_init = posterior_mu_init
@@ -448,7 +460,7 @@ def init_parameters(self):
             self.prior_bias_mu.data.fill_(self.prior_mean)
             self.prior_bias_sigma.data.fill_(self.prior_variance)
 
-    def forward(self, x):
+    def forward(self, x, return_kl=True):
 
         # linear outputs
         outputs = F.conv3d(x,
@@ -489,8 +501,11 @@ def forward(self, x):
                                      dilation=self.dilation,
                                      groups=self.groups) * sign_output
 
+        self.kl = kl
         # returning outputs + perturbations
-        return outputs + perturbed_outputs, kl
+        if return_kl:
+            return outputs + perturbed_outputs, kl
+        return outputs + perturbed_outputs
 
 
 class ConvTranspose1dFlipout(BaseVariationalLayer_):
@@ -537,6 +552,8 @@ def __init__(self,
         self.groups = groups
         self.bias = bias
 
+        self.kl = 0
+
         self.prior_mean = prior_mean
         self.prior_variance = prior_variance
         self.posterior_mu_init = posterior_mu_init
@@ -593,7 +610,7 @@ def init_parameters(self):
             self.prior_bias_mu.data.fill_(self.prior_mean)
             self.prior_bias_sigma.data.fill_(self.prior_variance)
 
-    def forward(self, x):
+    def forward(self, x, return_kl=True):
 
         # linear outputs
         outputs = F.conv_transpose1d(x,
@@ -635,8 +652,11 @@ def forward(self, x):
             dilation=self.dilation,
             groups=self.groups) * sign_output
 
+        self.kl = kl
         # returning outputs + perturbations
-        return outputs + perturbed_outputs, kl
+        if return_kl:
+            return outputs + perturbed_outputs, kl
+        return outputs + perturbed_outputs
 
 
 class ConvTranspose2dFlipout(BaseVariationalLayer_):
@@ -683,6 +703,8 @@ def __init__(self,
         self.groups = groups
         self.bias = bias
 
+        self.kl = 0
+
         self.prior_mean = prior_mean
         self.prior_variance = prior_variance
         self.posterior_mu_init = posterior_mu_init
@@ -743,7 +765,7 @@ def init_parameters(self):
             self.prior_bias_mu.data.fill_(self.prior_mean)
             self.prior_bias_sigma.data.fill_(self.prior_variance)
 
-    def forward(self, x):
+    def forward(self, x, return_kl=True):
 
         # linear outputs
         outputs = F.conv_transpose2d(x,
@@ -785,8 +807,11 @@ def forward(self, x):
             dilation=self.dilation,
             groups=self.groups) * sign_output
 
+        self.kl = kl
         # returning outputs + perturbations
-        return outputs + perturbed_outputs, kl
+        if return_kl:
+            return outputs + perturbed_outputs, kl
+        return outputs + perturbed_outputs
 
 
 class ConvTranspose3dFlipout(BaseVariationalLayer_):
@@ -838,6 +863,8 @@ def __init__(self,
         self.posterior_rho_init = posterior_rho_init
         self.bias = bias
 
+        self.kl = 0
+
         self.mu_kernel = nn.Parameter(
             torch.Tensor(in_channels, out_channels // groups, kernel_size,
                          kernel_size, kernel_size))
@@ -893,7 +920,7 @@ def init_parameters(self):
             self.prior_bias_mu.data.fill_(self.prior_mean)
             self.prior_bias_sigma.data.fill_(self.prior_variance)
 
-    def forward(self, x):
+    def forward(self, x, return_kl=True):
 
         # linear outputs
         outputs = F.conv_transpose3d(x,
@@ -935,5 +962,8 @@ def forward(self, x):
             dilation=self.dilation,
             groups=self.groups) * sign_output
 
+        self.kl = kl
         # returning outputs + perturbations
-        return outputs + perturbed_outputs, kl
+        if return_kl:
+            return outputs + perturbed_outputs, kl
+        return outputs + perturbed_outputs

bayesian_torch/layers/flipout_layers/linear_flipout.py

@@ -90,6 +90,8 @@ def __init__(self,
                              torch.Tensor(out_features, in_features),
                              persistent=False)
 
+        self.kl = 0
+
         if bias:
             self.mu_bias = nn.Parameter(torch.Tensor(out_features))
             self.rho_bias = nn.Parameter(torch.Tensor(out_features))
@@ -123,7 +125,7 @@ def init_parameters(self):
             self.mu_bias.data.normal_(mean=self.posterior_mu_init, std=0.1)
             self.rho_bias.data.normal_(mean=self.posterior_rho_init, std=0.1)
 
-    def forward(self, x):
+    def forward(self, x, return_kl=True):
         # sampling delta_W
         sigma_weight = torch.log1p(torch.exp(self.rho_weight))
         delta_weight = (sigma_weight * self.eps_weight.data.normal_())
@@ -148,5 +150,9 @@ def forward(self, x):
         perturbed_outputs = F.linear(x * sign_input, delta_weight,
                                      bias) * sign_output
 
+        self.kl = kl
+
         # returning outputs + perturbations
-        return outputs + perturbed_outputs, kl
+        if return_kl:
+            return outputs + perturbed_outputs, kl
+        return outputs + perturbed_outputs

bayesian_torch/layers/flipout_layers/rnn_flipout.py

@@ -76,6 +76,8 @@ def __init__(self,
         self.posterior_rho_init = posterior_rho_init,  # variance of weight --> sigma = log (1 + exp(rho))
         self.bias = bias
 
+        self.kl = 0
+
         self.ih = LinearFlipout(prior_mean=prior_mean,
                                 prior_variance=prior_variance,
                                 posterior_mu_init=posterior_mu_init,
@@ -92,7 +94,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()
 
@@ -137,4 +139,7 @@ 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)