bayesian segnet

Author: Krzysztof Maciej Lis

.gitignore

@@ -0,0 +1,2 @@
+ __pycache__/
+ models/seg_net_bayes.py

models/__init__.py

@@ -5,4 +5,5 @@
 from .gcn import *
 from .psp_net import *
 from .seg_net import *
+from .seg_net_bayes import *
 from .u_net import *

models/psp_net.py

@@ -9,7 +9,7 @@
 
 class _PyramidPoolingModule(nn.Module):
     def __init__(self, in_dim, reduction_dim, setting):
-        super(_PyramidPoolingModule, self).__init__()
+        super().__init__()
         self.features = []
         for s in setting:
             self.features.append(nn.Sequential(
@@ -24,14 +24,14 @@ def forward(self, x):
         x_size = x.size()
         out = [x]
         for f in self.features:
-            out.append(F.upsample(f(x), x_size[2:], mode='bilinear'))
+            out.append(F.interpolate(f(x), x_size[2:], mode='bilinear'))
         out = torch.cat(out, 1)
         return out
 
 
 class PSPNet(nn.Module):
     def __init__(self, num_classes, pretrained=True, use_aux=True):
-        super(PSPNet, self).__init__()
+        super().__init__()
         self.use_aux = use_aux
         resnet = models.resnet101(pretrained=pretrained)
         self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool)
@@ -63,7 +63,8 @@ def __init__(self, num_classes, pretrained=True, use_aux=True):
 
         initialize_weights(self.ppm, self.final)
 
-    def forward(self, x):
+    def forward(self, image):
+        x = image
         x_size = x.size()
         x = self.layer0(x)
         x = self.layer1(x)
@@ -75,8 +76,8 @@ def forward(self, x):
         x = self.ppm(x)
         x = self.final(x)
         if self.training and self.use_aux:
-            return F.upsample(x, x_size[2:], mode='bilinear'), F.upsample(aux, x_size[2:], mode='bilinear')
-        return F.upsample(x, x_size[2:], mode='bilinear')
+            return F.interpolate(x, x_size[2:], mode='bilinear'), F.interpolate(aux, x_size[2:], mode='bilinear')
+        return F.interpolate(x, x_size[2:], mode='bilinear')
 
 
 # just a try, not recommend to use

models/seg_net.py

@@ -8,7 +8,7 @@
 class _DecoderBlock(nn.Module):
     def __init__(self, in_channels, out_channels, num_conv_layers):
         super(_DecoderBlock, self).__init__()
-        middle_channels = in_channels / 2
+        middle_channels = in_channels // 2
         layers = [
             nn.ConvTranspose2d(in_channels, in_channels, kernel_size=2, stride=2),
             nn.Conv2d(in_channels, middle_channels, kernel_size=3, padding=1),

models/seg_net_bayes.py

@@ -0,0 +1,99 @@
+
+import torch
+from torch import nn
+from torchvision import models
+
+from ..utils import initialize_weights
+from .seg_net import _DecoderBlock, SegNet
+
+class SegNetBayes(SegNet):
+	def __init__(self, num_classes, dropout_p=0.5, pretrained=True, num_samples=16, min_batch_size=4):
+		"""
+		@param num_samples: number of samples for the Monte-Carlo simulation,
+			how many times to run with random dropout
+		"""
+		super().__init__(num_classes=num_classes, pretrained=pretrained)
+
+		self.drop = nn.Dropout2d(p=dropout_p, inplace=False)
+		self.num_samples = num_samples
+		self.min_batch_size = min_batch_size
+
+	def forward(self, x):
+		enc1 = self.enc1(x)
+		#print('enc1', enc1.shape)
+
+		enc2 = self.enc2(enc1)
+		#print('enc2', enc2.shape)
+
+		enc3 = self.enc3(enc2)
+		#print('enc3', enc3.shape)
+		enc3 = self.drop(enc3)
+		#print('enc3d', enc3.shape)
+
+		enc4 = self.enc4(enc3)
+		#print('enc4', enc4.shape)
+		enc4 = self.drop(enc4)
+		#print('enc4d', enc4.shape)
+
+		enc5 = self.enc5(enc4)
+		#print('enc5', enc5.shape)
+		enc5 = self.drop(enc5)
+		#print('enc5d', enc5.shape)
+
+		dec5 = self.dec5(enc5)
+		#print('dec5', dec5.shape)
+		dec5 = self.drop(dec5)
+		#print('dec5d', dec5.shape)
+
+		dec4 = self.dec4(torch.cat([enc4, dec5], 1))
+		#print('dec4', dec4.shape)
+		dec4 = self.drop(dec4)
+		#print('dec4d', dec4.shape)
+
+		dec3 = self.dec3(torch.cat([enc3, dec4], 1))
+		dec3 = self.drop(dec3)
+
+		dec2 = self.dec2(torch.cat([enc2, dec3], 1))
+		dec1 = self.dec1(torch.cat([enc1, dec2], 1))
+		return dec1
+
+	def forward_multisample(self, x, num_samples=None):
+		# dropout must be active
+		backup_train_mode = self.drop.training
+		self.drop.train()
+
+		softmax = torch.nn.Softmax2d()
+
+		num_samples = num_samples if num_samples else self.num_samples
+
+		results = [softmax(self.forward(x)).data.cpu() for i in range(num_samples)]
+
+		preds = torch.stack(results).cuda()
+		avg = torch.mean(preds, 0)
+		var = torch.var(preds, 0)
+		del preds
+
+		# restore mode
+		self.drop.train(backup_train_mode)
+
+		return dict(
+			mean = avg,
+			var = var,
+		)
+
+	#def sample(self, x, num_samples, batch_size):
+		#infer desired batch size from input shape
+		#we will divide a num_samples into batches
+		#num_frames = x.shape[0]
+		#batch_size = max(num_frames, self.min_batch_size)
+
+		#for sample_idx in range(num_samples):
+			#pred =
+
+
+		#for fr_idx in range(num_frames):
+			#x_single = x[fr_idx:fr_idx+1, :, :, :]
+			#self.sample(x_single, num_samples, batch_size)
+
+
+

utils/misc.py

@@ -17,7 +17,7 @@ def initialize_weights(*models):
     for model in models:
         for module in model.modules():
             if isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear):
-                nn.init.kaiming_normal(module.weight)
+                nn.init.kaiming_normal_(module.weight)
                 if module.bias is not None:
                     module.bias.data.zero_()
             elif isinstance(module, nn.BatchNorm2d):