Merge pull request #628 from lukeyeager/layer-stages

Use layer stages for all-in-one nets
NVIDIA · Mar 15, 2016 · 4bd5b9e · 4bd5b9e
2 parents 41996c8 + a5eba04
commit 4bd5b9e
Show file tree

Hide file tree

Showing 8 changed files with 503 additions and 181 deletions.
diff --git a/digits/model/images/classification/test_views.py b/digits/model/images/classification/test_views.py
@@ -57,16 +57,22 @@ class BaseViewsTest(digits.test_views.BaseViewsTest):
     bottom: "output"
     bottom: "label"
     top: "loss"
+    exclude { stage: "deploy" }
 }
 layer {
     name: "accuracy"
     type: "Accuracy"
     bottom: "output"
     bottom: "label"
     top: "accuracy"
-    include {
-        phase: TEST
-    }
+    include { stage: "val" }
+}
+layer {
+    name: "softmax"
+    type: "Softmax"
+    bottom: "output"
+    top: "softmax"
+    include { stage: "deploy" }
 }
 """
 
@@ -431,6 +437,14 @@ def test_bad_network_definition(self):
                     bottom: "output"
                     bottom: "label"
                     top: "loss"
+                    exclude { stage: "deploy" }
+                }
+                layer {
+                    name: "softmax"
+                    type: "Softmax"
+                    bottom: "output"
+                    top: "softmax"
+                    include { stage: "deploy" }
                 }
                 """
         elif self.FRAMEWORK == 'torch':
@@ -839,16 +853,22 @@ class BaseTestCreatedCropInNetwork(BaseTestCreated):
     bottom: "output"
     bottom: "label"
     top: "loss"
+    exclude { stage: "deploy" }
 }
 layer {
     name: "accuracy"
     type: "Accuracy"
     bottom: "output"
     bottom: "label"
     top: "accuracy"
-    include {
-        phase: TEST
-    }
+    include { stage: "val" }
+}
+layer {
+    name: "softmax"
+    type: "Softmax"
+    bottom: "output"
+    top: "softmax"
+    include { stage: "deploy" }
 }
 """
     TORCH_NETWORK = \
@@ -1006,16 +1026,15 @@ class TestPythonLayer(BaseViewsTestWithDataset):
     bottom: "output"
     bottom: "label"
     top: "loss"
+    exclude { stage: "deploy" }
 }
 layer {
     name: "accuracy"
     type: "Accuracy"
     bottom: "output"
     bottom: "label"
     top: "accuracy"
-    include {
-        phase: TEST
-    }
+    include { stage: "val" }
 }
 layer {
     name: "py_test"
@@ -1027,7 +1046,13 @@ class TestPythonLayer(BaseViewsTestWithDataset):
         layer: "PythonLayer"
     }
 }
-
+layer {
+    name: "softmax"
+    type: "Softmax"
+    bottom: "output"
+    top: "softmax"
+    include { stage: "deploy" }
+}
 """
     def write_python_layer_script(self, filename):
         with open(filename, 'w') as f:

diff --git a/digits/model/images/generic/test_views.py b/digits/model/images/generic/test_views.py
@@ -66,11 +66,12 @@ class BaseViewsTest(digits.test_views.BaseViewsTest):
   }
 }
 layer {
-  name: "train_loss"
+  name: "loss"
   type: "EuclideanLoss"
   bottom: "output"
   bottom: "label"
   top: "loss"
+  exclude { stage: "deploy" }
 }
 """
 
@@ -711,11 +712,12 @@ class BaseTestCreatedCropInNetwork(BaseTestCreated):
   }
 }
 layer {
-  name: "train_loss"
+  name: "loss"
   type: "EuclideanLoss"
   bottom: "output"
   bottom: "label"
   top: "loss"
+  exclude { stage: "deploy" }
 }
 """