Merge pull request #43 from kazuto1011/develop

torch.hub support
kazuto1011 · Dec 20, 2018 · 99312a1 · 99312a1
2 parents 6b39584 + 863c444
commit 99312a1
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 6 deletions.
diff --git a/README.md b/README.md
@@ -16,8 +16,6 @@ For anaconda users:
 
 ```sh
 conda env create --file config/conda_env.yaml
-conda activate deeplab-pytorch
-conda install pytorch torchvision -c pytorch  # depends on your environment
 ```
 
 * python 2.7/3.6
@@ -211,6 +209,17 @@ python livedemo.py --config config/cocostuff164k.yaml \
                    --camera-id <CAMERA ID>
 ```
 
+### torch.hub
+
+```python
+import torch.hub
+
+model = torch.hub.load(
+    "kazuto1011/deeplab-pytorch", "deeplabv2_resnet101", n_classes=182
+)
+model.load_state_dict(torch.load("cocostuff164k_iter100k.pth"))
+```
+
 ## References
 
 1. [DeepLab: Semantic Image Segmentation with Deep Convolutional Nets, Atrous Convolution, and Fully Connected CRFs](https://arxiv.org/abs/1606.00915)<br>

diff --git a/config/conda_env.yaml b/config/conda_env.yaml
@@ -4,16 +4,17 @@ channels:
   - conda-forge
   - defaults
 dependencies:
-  # - pytorch=0.4.1
-  # - torchvision=0.2.1
+  - pytorch
+  - torchvision
+  # - cuda92
   - h5py
   - scipy
   - matplotlib
   - pyyaml
   - click
   - tqdm
-  - clang
-  - clangxx
+  # - clang
+  # - clangxx
   - pydensecrf
   - pip:
     - torchnet==0.0.2

diff --git a/hubconf.py b/hubconf.py
@@ -0,0 +1,21 @@
+#!/usr/bin/env python
+# coding: utf-8
+#
+# Author: Kazuto Nakashima
+# URL:    https://kazuto1011.github.io
+# Date:   20 December 2018
+
+
+def deeplabv2_resnet101(**kwargs):
+    """
+    DeepLab v2 model with ResNet-101 backbone
+    n_classes (int): the number of classes
+    """
+
+    from libs.models.deeplabv2 import DeepLabV2
+    from libs.models.msc import MSC
+
+    base = DeepLabV2(n_blocks=[3, 4, 23, 3], pyramids=[6, 12, 18, 24], **kwargs)
+    model = MSC(scale=base, pyramids=[0.5, 0.75])
+
+    return model