Change the dependencies

Author: kazuto1011

configs/conda_env.yaml

@@ -1,26 +1,20 @@
 name: deeplab-pytorch
-channels:
-  - pytorch
-  - conda-forge
-  - defaults
 dependencies:
-  # - clang # For MacOS
-  # - clangxx # For MacOS
   - click
-  - cudatoolkit=10.0
-  - cython
+  - conda-forge::pydensecrf
+  - cudatoolkit=10.2
   - matplotlib
-  - pydensecrf
-  - pytorch
+  - python=3.6
+  - pytorch::pytorch>1.2.0
+  - pytorch::torchvision
   - pyyaml
   - scipy
-  - torchvision
   - tqdm
   - pip:
     - addict
     - black
     - joblib
+    - omegaconf
     - opencv-python
-    - tensorboardX
     - tensorflow
     - torchnet

demo.py

@@ -16,8 +16,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import yaml
-from addict import Dict
+from omegaconf import OmegaConf
 
 from libs.models import *
 from libs.utils import DenseCRF
@@ -139,7 +138,7 @@ def single(config_path, model_path, image_path, cuda, crf):
     """
 
     # Setup
-    CONFIG = Dict(yaml.load(config_path))
+    CONFIG = OmegaConf.load(config_path)
     device = get_device(cuda)
     torch.set_grad_enabled(False)
 
@@ -207,7 +206,7 @@ def live(config_path, model_path, cuda, crf, camera_id):
     """
 
     # Setup
-    CONFIG = Dict(yaml.load(config_path))
+    CONFIG = OmegaConf.load(config_path)
     device = get_device(cuda)
     torch.set_grad_enabled(False)
     torch.backends.cudnn.benchmark = True

libs/models/resnet.py

@@ -60,7 +60,7 @@ def __init__(self, in_ch, out_ch, stride, dilation, downsample):
         self.shortcut = (
             _ConvBnReLU(in_ch, out_ch, 1, stride, 0, 1, False)
             if downsample
-            else lambda x: x  # identity
+            else nn.Identity()
         )
 
     def forward(self, x):
@@ -110,11 +110,6 @@ def __init__(self, out_ch):
         self.add_module("pool", nn.MaxPool2d(3, 2, 1, ceil_mode=True))
 
 
-class _Flatten(nn.Module):
-    def forward(self, x):
-        return x.view(x.size(0), -1)
-
-
 class ResNet(nn.Sequential):
     def __init__(self, n_classes, n_blocks):
         super(ResNet, self).__init__()
@@ -125,7 +120,7 @@ def __init__(self, n_classes, n_blocks):
         self.add_module("layer4", _ResLayer(n_blocks[2], ch[3], ch[4], 2, 1))
         self.add_module("layer5", _ResLayer(n_blocks[3], ch[4], ch[5], 2, 1))
         self.add_module("pool5", nn.AdaptiveAvgPool2d(1))
-        self.add_module("flatten", _Flatten())
+        self.add_module("flatten", nn.Flatten())
         self.add_module("fc", nn.Linear(ch[5], n_classes))
 
 

main.py

@@ -17,15 +17,14 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import yaml
-from addict import Dict
+from omegaconf import OmegaConf
 from PIL import Image
-from tensorboardX import SummaryWriter
+from torch.utils.tensorboard import SummaryWriter
 from torchnet.meter import MovingAverageValueMeter
 from tqdm import tqdm
 
 from libs.datasets import get_dataset
-from libs.models import *
+from libs.models import DeepLabV2_ResNet101_MSC
 from libs.utils import DenseCRF, PolynomialLR, scores
 
 
@@ -110,7 +109,7 @@ def train(config_path, cuda):
     """
 
     # Configuration
-    CONFIG = Dict(yaml.load(config_path))
+    CONFIG = OmegaConf.load(config_path)
     device = get_device(cuda)
     torch.backends.cudnn.benchmark = True
 
@@ -144,7 +143,7 @@ def train(config_path, cuda):
     ), 'Currently support only "DeepLabV2_ResNet101_MSC"'
 
     # Model setup
-    model = eval(CONFIG.MODEL.NAME)(n_classes=CONFIG.DATASET.N_CLASSES)
+    model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.DATASET.N_CLASSES)
     state_dict = torch.load(CONFIG.MODEL.INIT_MODEL)
     print("    Init:", CONFIG.MODEL.INIT_MODEL)
     for m in model.base.state_dict().keys():
@@ -236,7 +235,7 @@ def train(config_path, cuda):
                 labels_ = resize_labels(labels, size=(H, W))
                 iter_loss += criterion(logit, labels_.to(device))
 
-            # Propagate backward (just compute gradients wrt the loss)
+            # Propagate backward (just compute gradients)
             iter_loss /= CONFIG.SOLVER.ITER_SIZE
             iter_loss.backward()
 
@@ -308,7 +307,7 @@ def test(config_path, model_path, cuda):
     """
 
     # Configuration
-    CONFIG = Dict(yaml.load(config_path))
+    CONFIG = OmegaConf.load(config_path)
     device = get_device(cuda)
     torch.set_grad_enabled(False)
 
@@ -417,7 +416,7 @@ def crf(config_path, n_jobs):
     """
 
     # Configuration
-    CONFIG = Dict(yaml.load(config_path))
+    CONFIG = OmegaConf.load(config_path)
     torch.set_grad_enabled(False)
     print("# jobs:", n_jobs)