Merge pull request #2 from y0ast/main

Fix KL computation

Author: ranganathkrishnan

bayesian_torch/models/bayesian/resnet_flipout.py

@@ -132,13 +132,13 @@ def forward(self, x):
         out = F.relu(out)
         for l in self.layer1:
             out, kl = l(out)
-        kl_sum += kl
+            kl_sum += kl
         for l in self.layer2:
             out, kl = l(out)
-        kl_sum += kl
+            kl_sum += kl
         for l in self.layer3:
             out, kl = l(out)
-        kl_sum += kl
+            kl_sum += kl
 
         out = F.avg_pool2d(out, out.size()[3])
         out = out.view(out.size(0), -1)

bayesian_torch/models/bayesian/resnet_flipout_large.py

@@ -215,36 +215,20 @@ def forward(self, x):
         x = self.maxpool(x)
 
         for layer in self.layer1:
-            if 'Flipout' in str(layer):
-                x, kl = layer(x)
-                if kl is None:
-                    kl_sum += kl
-            else:
-                x = layer(x)
+            x, kl = layer(x)
+            kl_sum += kl
 
         for layer in self.layer2:
-            if 'Flipout' in str(layer):
-                x, kl = layer(x)
-                if kl is None:
-                    kl_sum += kl
-            else:
-                x = layer(x)
+            x, kl = layer(x)
+            kl_sum += kl
 
         for layer in self.layer3:
-            if 'Flipout' in str(layer):
-                x, kl = layer(x)
-                if kl is None:
-                    kl_sum += kl
-            else:
-                x = layer(x)
+            x, kl = layer(x)
+            kl_sum += kl
 
         for layer in self.layer4:
-            if 'Flipout' in str(layer):
-                x, kl = layer(x)
-                if kl is None:
-                    kl_sum += kl
-            else:
-                x = layer(x)
+            x, kl = layer(x)
+            kl_sum += kl
 
         x = self.avgpool(x)
         x = x.view(x.size(0), -1)

bayesian_torch/models/bayesian/resnet_variational.py

@@ -152,13 +152,13 @@ def forward(self, x):
         out = F.relu(out)
         for l in self.layer1:
             out, kl = l(out)
-        kl_sum += kl
+            kl_sum += kl
         for l in self.layer2:
             out, kl = l(out)
-        kl_sum += kl
+            kl_sum += kl
         for l in self.layer3:
             out, kl = l(out)
-        kl_sum += kl
+            kl_sum += kl
 
         out = F.avg_pool2d(out, out.size()[3])
         out = out.view(out.size(0), -1)

bayesian_torch/models/bayesian/resnet_variational_large.py

@@ -220,36 +220,20 @@ def forward(self, x):
         x = self.maxpool(x)
 
         for layer in self.layer1:
-            if 'Reparameterization' in str(layer):
-                x, kl = layer(x)
-                if kl is None:
-                    kl_sum += kl
-            else:
-                x = layer(x)
+            x, kl = layer(x)
+            kl_sum += kl
 
         for layer in self.layer2:
-            if 'Reparameterization' in str(layer):
-                x, kl = layer(x)
-                if kl is None:
-                    kl_sum += kl
-            else:
-                x = layer(x)
+            x, kl = layer(x)
+            kl_sum += kl
 
         for layer in self.layer3:
-            if 'Reparameterization' in str(layer):
-                x, kl = layer(x)
-                if kl is None:
-                    kl_sum += kl
-            else:
-                x = layer(x)
+            x, kl = layer(x)
+            kl_sum += kl
 
         for layer in self.layer4:
-            if 'Reparameterization' in str(layer):
-                x, kl = layer(x)
-                if kl is None:
-                    kl_sum += kl
-            else:
-                x = layer(x)
+            x, kl = layer(x)
+            kl_sum += kl
 
         x = self.avgpool(x)
         x = x.view(x.size(0), -1)