PyTorch 1.0.1 and rasterio, fiona (#72)

* Edits for PyTorch 1.0.x

* interpolate replaces upsample

Author: ymoisan

Diff for: inference.py

@@ -186,7 +186,7 @@ def sem_seg_inference(bucket, model, image, overlay):
                     col_to = col + chunk_size - overlay
 
                     useful_sem_seg = segmentation[overlay:chunk_size - overlay, overlay:chunk_size - overlay]
-                    output_np[row_from:row_to, col_from:col_to, 0] = useful_sem_seg
+                    output_np[row_from:row_to, col_from:col_to, 0] = useful_sem_seg.cpu()
 
             # Resize the output array to the size of the input image and write it
             output_np = output_np[overlay:h + overlay, overlay:w + overlay]

Diff for: metrics.py

@@ -53,7 +53,7 @@ def report_classification(pred, label, batch_size, metrics_dict):
     """Computes precision, recall and f-score for each class and average of all classes.
     http://scikit-learn.org/stable/modules/generated/sklearn.metrics.classification_report.html
     """
-    class_report = classification_report(label, pred, output_dict=True)
+    class_report = classification_report(label.cpu(), pred.cpu(), output_dict=True)
 
     class_score = {}
     for key, value in class_report.items():
@@ -75,6 +75,6 @@ def iou(pred, target, batch_size, metrics_dict):
     """Calculate the intersection over union (or Jaccard index) between two datasets.
     The Jaccard distance (or dissimilarity) would be 1-iou.
     """
-    iou = jaccard_similarity_score(target, pred, normalize=True)
+    iou = jaccard_similarity_score(target.cpu(), pred.cpu(), normalize=True)
     metrics_dict['iou'].update(iou, batch_size)
     return metrics_dict

Diff for: models/checkpointed_unet.py

@@ -1,4 +1,4 @@
-import torch
+import torch, utils
 # import torch should be first. Unclear issue, mentioned here: https://github.com/pytorch/pytorch/issues/2083
 from torch import nn
 from torch.utils.checkpoint import checkpoint_sequential
@@ -44,15 +44,14 @@ def forward(self, input_data):
         modules = get_modules(self.encoding_block)
         return checkpoint_sequential(modules, segments, input_data)
 
-
 class DecodingBlock(nn.Module):
     """Module in the decoding section of the UNet"""
 
     def __init__(self, in_size, out_size, batch_norm=False, upsampling=True):
         super().__init__()
         up_modules = []
         if upsampling:
-            self.up = nn.Sequential(nn.Upsample(mode='bilinear', scale_factor=2),
+            self.up = nn.Sequential(utils.Interpolate(mode='bilinear', scale_factor=2),
                                     nn.Conv2d(in_size, out_size, kernel_size=1))
             self.upsampling = True
         else:
@@ -71,7 +70,7 @@ def forward(self, input1, input2):
             output2 = checkpoint_sequential(self.up_modules, segments, input2)
         else:
             output2 = self.up(input2)
-        output1 = nn.functional.upsample(input1, output2.size()[2:], mode='bilinear', align_corners=True)
+        output1 = nn.functional.interpolate(input1, output2.size()[2:], mode='bilinear', align_corners=True)
         return checkpoint_sequential(self.conv_modules, segments, torch.cat([output1, output2], 1))
 
 
@@ -118,7 +117,7 @@ def forward(self, input_data):
         decode2 = self.decode2(conv2, decode3)
         decode1 = self.decode1(conv1, decode2)
 
-        final = nn.functional.upsample(self.final(decode1), input_data.size()[2:], mode='bilinear')
+        final = nn.functional.interpolate(self.final(decode1), input_data.size()[2:], mode='bilinear')
         return final
 
 
@@ -157,7 +156,7 @@ def forward(self, input_data):
         decode2 = self.decode2(conv2, decode3)
         decode1 = self.decode1(conv1, decode2)
 
-        final = nn.functional.upsample(self.final(decode1), input_data.size()[2:], mode='bilinear', align_corners=True)
+        final = nn.functional.interpolate(self.final(decode1), input_data.size()[2:], mode='bilinear', align_corners=True)
         return final
 
 

Diff for: models/unet.py

@@ -1,4 +1,4 @@
-import torch
+import torch, utils
 from torch import nn
 
 
@@ -48,7 +48,7 @@ class DecodingBlock(nn.Module):
     def __init__(self, in_size, out_size, batch_norm=False, upsampling=True):
         super().__init__()
         if upsampling:
-            self.up = nn.Sequential(nn.Upsample(mode='bilinear', scale_factor=2),
+            self.up = nn.Sequential(utils.Interpolate(mode='bilinear', scale_factor=2),
                                     nn.Conv2d(in_size, out_size, kernel_size=1))
         else:
             self.up = nn.ConvTranspose2d(in_size, out_size, kernel_size=2, stride=2)
@@ -57,7 +57,7 @@ def __init__(self, in_size, out_size, batch_norm=False, upsampling=True):
 
     def forward(self, input1, input2):
         output2 = self.up(input2)
-        output1 = nn.functional.upsample(input1, output2.size()[2:], mode='bilinear', align_corners=True)
+        output1 = nn.functional.interpolate(input1, output2.size()[2:], mode='bilinear', align_corners=True)
         return self.conv(torch.cat([output1, output2], 1))
 
 
@@ -102,7 +102,7 @@ def forward(self, input_data):
         decode2 = self.decode2(conv2, decode3)
         decode1 = self.decode1(conv1, decode2)
 
-        final = nn.functional.upsample(self.final(decode1), input_data.size()[2:], mode='bilinear')
+        final = nn.functional.interpolate(self.final(decode1), input_data.size()[2:], mode='bilinear')
 
         return final
 
@@ -142,6 +142,6 @@ def forward(self, input_data):
         decode2 = self.decode2(conv2, decode3)
         decode1 = self.decode1(conv1, decode2)
 
-        final = nn.functional.upsample(self.final(decode1), input_data.size()[2:], mode='bilinear', align_corners=True)
+        final = nn.functional.interpolate(self.final(decode1), input_data.size()[2:], mode='bilinear', align_corners=True)
 
         return final

Diff for: utils.py

@@ -13,6 +13,17 @@
     warnings.warn('The boto3 library counldn\'t be imported. Ignore if not using AWS s3 buckets', ImportWarning)
     pass
 
+class Interpolate(torch.nn.Module):
+    def __init__(self, mode, scale_factor):
+        super(Interpolate, self).__init__()
+        self.interp = torch.nn.functional.interpolate
+        self.scale_factor = scale_factor
+        self.mode = mode
+
+    def forward(self, x):
+        x = self.interp(x, scale_factor=self.scale_factor, mode=self.mode, align_corners=False)
+        return x
+
 
 def create_or_empty_folder(folder):
     """Empty an existing folder or create it if it doesn't exist.