include kl_loss() function in Convolutional flipout layers,

to compute kl when 'return_kl' flag is set to False.
Fix for issue#12.

Signed-off-by: Ranganath Krishnan <[email protected]>

Author: ranganathkrishnan

bayesian_torch/layers/flipout_layers/conv_flipout.py

@@ -152,6 +152,14 @@ def init_parameters(self):
             self.prior_bias_mu.data.fill_(self.prior_mean)
             self.prior_bias_sigma.data.fill_(self.prior_variance)
 
+    def kl_loss(self):
+        sigma_weight = torch.log1p(torch.exp(self.rho_kernel))
+        kl = self.kl_div(self.mu_kernel, sigma_weight, self.prior_weight_mu, self.prior_weight_sigma)
+        if self.bias:
+           sigma_bias = torch.log1p(torch.exp(self.rho_bias))
+           kl += self.kl_div(self.mu_bias, sigma_bias, self.prior_bias_mu, self.prior_bias_sigma)
+        return kl
+
     def forward(self, x, return_kl=True):
 
         if self.dnn_to_bnn_flag:
@@ -311,6 +319,14 @@ def init_parameters(self):
             self.prior_bias_mu.data.fill_(self.prior_mean)
             self.prior_bias_sigma.data.fill_(self.prior_variance)
 
+    def kl_loss(self):
+        sigma_weight = torch.log1p(torch.exp(self.rho_kernel))
+        kl = self.kl_div(self.mu_kernel, sigma_weight, self.prior_weight_mu, self.prior_weight_sigma)
+        if self.bias:
+           sigma_bias = torch.log1p(torch.exp(self.rho_bias))
+           kl += self.kl_div(self.mu_bias, sigma_bias, self.prior_bias_mu, self.prior_bias_sigma)
+        return kl
+
     def forward(self, x, return_kl=True):
 
         if self.dnn_to_bnn_flag:
@@ -469,6 +485,14 @@ def init_parameters(self):
             self.prior_bias_mu.data.fill_(self.prior_mean)
             self.prior_bias_sigma.data.fill_(self.prior_variance)
 
+    def kl_loss(self):
+        sigma_weight = torch.log1p(torch.exp(self.rho_kernel))
+        kl = self.kl_div(self.mu_kernel, sigma_weight, self.prior_weight_mu, self.prior_weight_sigma)
+        if self.bias:
+           sigma_bias = torch.log1p(torch.exp(self.rho_bias))
+           kl += self.kl_div(self.mu_bias, sigma_bias, self.prior_bias_mu, self.prior_bias_sigma)
+        return kl
+
     def forward(self, x, return_kl=True):
 
         if self.dnn_to_bnn_flag:
@@ -624,6 +648,14 @@ def init_parameters(self):
             self.prior_bias_mu.data.fill_(self.prior_mean)
             self.prior_bias_sigma.data.fill_(self.prior_variance)
 
+    def kl_loss(self):
+        sigma_weight = torch.log1p(torch.exp(self.rho_kernel))
+        kl = self.kl_div(self.mu_kernel, sigma_weight, self.prior_weight_mu, self.prior_weight_sigma)
+        if self.bias:
+           sigma_bias = torch.log1p(torch.exp(self.rho_bias))
+           kl += self.kl_div(self.mu_bias, sigma_bias, self.prior_bias_mu, self.prior_bias_sigma)
+        return kl
+
     def forward(self, x, return_kl=True):
 
         if self.dnn_to_bnn_flag:
@@ -784,6 +816,14 @@ def init_parameters(self):
             self.prior_bias_mu.data.fill_(self.prior_mean)
             self.prior_bias_sigma.data.fill_(self.prior_variance)
 
+    def kl_loss(self):
+        sigma_weight = torch.log1p(torch.exp(self.rho_kernel))
+        kl = self.kl_div(self.mu_kernel, sigma_weight, self.prior_weight_mu, self.prior_weight_sigma)
+        if self.bias:
+           sigma_bias = torch.log1p(torch.exp(self.rho_bias))
+           kl += self.kl_div(self.mu_bias, sigma_bias, self.prior_bias_mu, self.prior_bias_sigma)
+        return kl
+
     def forward(self, x, return_kl=True):
 
         if self.dnn_to_bnn_flag:
@@ -944,6 +984,14 @@ def init_parameters(self):
             self.prior_bias_mu.data.fill_(self.prior_mean)
             self.prior_bias_sigma.data.fill_(self.prior_variance)
 
+    def kl_loss(self):
+        sigma_weight = torch.log1p(torch.exp(self.rho_kernel))
+        kl = self.kl_div(self.mu_kernel, sigma_weight, self.prior_weight_mu, self.prior_weight_sigma)
+        if self.bias:
+           sigma_bias = torch.log1p(torch.exp(self.rho_bias))
+           kl += self.kl_div(self.mu_bias, sigma_bias, self.prior_bias_mu, self.prior_bias_sigma)
+        return kl
+
     def forward(self, x, return_kl=True):
 
         if self.dnn_to_bnn_flag: