Repository: Lightning-AI/pytorch-lightning

Issue:
Fix segmentation example
# Before submitting

- [x] Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md)?
- [x] Did you make sure to update the docs?   
- [ ] Did you write any new necessary tests?  

## What does this PR do?
Fixes #874 (issue).

## PR review    
Anyone in the community is free to review the PR once the tests have passed. 

## Some comments
I have added a custom model (UNet) instead of the previous torchvision segmentation model as it was causing tests to fail (mostly because those torchvision models were introduced in version 0.3 and it will fail for previous versions).


Discussion:
Hello @akshaykvnit! Thanks for updating this PR. 









There are currently no PEP 8 issues detected in this Pull Request. Cheers! :beers: 

##### Comment last updated at 2020-02-17 06:20:52 UTC
I made the changes suggested by @Borda, but @williamFalcon already merged the previous commit :sweat_smile: 
> I made the changes suggested by @Borda, but @williamFalcon already merged the previous commit

@akshaykvnit could you pls do follow-up PR? kind of Deja Vu lol

Resolving patch:
diff --git a/pl_examples/full_examples/semantic_segmentation/models/unet/__init__.py b/pl_examples/full_examples/semantic_segmentation/models/unet/__init__.py
deleted file mode 100644
--- a/pl_examples/full_examples/semantic_segmentation/models/unet/__init__.py
+++ /dev/null
@@ -1,4 +0,0 @@
-# For relative imports to work in Python 3.6
-import os
-import sys
-sys.path.append(os.path.dirname(os.path.realpath(__file__)))
diff --git a/pl_examples/full_examples/semantic_segmentation/models/unet/model.py b/pl_examples/full_examples/semantic_segmentation/models/unet/model.py
--- a/pl_examples/full_examples/semantic_segmentation/models/unet/model.py
+++ b/pl_examples/full_examples/semantic_segmentation/models/unet/model.py
@@ -2,30 +2,33 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from parts import DoubleConv, Down, Up
+from models.unet.parts import DoubleConv, Down, Up
 
 
 class UNet(nn.Module):
     '''
     Architecture based on U-Net: Convolutional Networks for Biomedical Image Segmentation
     Link - https://arxiv.org/abs/1505.04597
+
+    Parameters:
+    num_classes (int) - Number of output classes required (default 19 for KITTI dataset)
+    bilinear (bool) - Whether to use bilinear interpolation or transposed
+    convolutions for upsampling.
     '''
     def __init__(self, num_classes=19, bilinear=False):
         super().__init__()
-        self.bilinear = bilinear
-        self.num_classes = num_classes
         self.layer1 = DoubleConv(3, 64)
         self.layer2 = Down(64, 128)
         self.layer3 = Down(128, 256)
         self.layer4 = Down(256, 512)
         self.layer5 = Down(512, 1024)
 
-        self.layer6 = Up(1024, 512, bilinear=self.bilinear)
-        self.layer7 = Up(512, 256, bilinear=self.bilinear)
-        self.layer8 = Up(256, 128, bilinear=self.bilinear)
-        self.layer9 = Up(128, 64, bilinear=self.bilinear)
+        self.layer6 = Up(1024, 512, bilinear=bilinear)
+        self.layer7 = Up(512, 256, bilinear=bilinear)
+        self.layer8 = Up(256, 128, bilinear=bilinear)
+        self.layer9 = Up(128, 64, bilinear=bilinear)
 
-        self.layer10 = nn.Conv2d(64, self.num_classes, kernel_size=1)
+        self.layer10 = nn.Conv2d(64, num_classes, kernel_size=1)
 
     def forward(self, x):
         x1 = self.layer1(x)


Test patch: