Add tests for ONNX pass scopes

lantiga · lantiga · commit 447b17d40eef · 2018-01-08T08:39:16.000+01:00
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -158,6 +158,47 @@ def f(x, y):
         trace, z = torch.jit.trace(f, (x, y), nderivs=0)
         self.assertExpectedTrace(trace)
 
+        class Net(nn.Module):
+            def forward(self, x):
+                return F.log_softmax(x, dim=0)
+
+        net = Net()
+        t = Variable(torch.ones(2), requires_grad=True)
+        trace, _ = torch.jit.trace(net, (t, ))
+        torch.onnx._optimize_trace(trace, False)
+        g = torch._C._jit_get_graph(trace)
+        for node in g.nodes():
+            self.assertTrue(node.scopeName() == 'Net')
+
+        class Net(nn.Module):
+
+            def __init__(self, num_classes=1000):
+                super(Net, self).__init__()
+                self.features = nn.Sequential(
+                    nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
+                    nn.ReLU(inplace=True),
+                    nn.MaxPool2d(kernel_size=3, stride=2),
+                )
+
+            def forward(self, x):
+                x = self.features(x)
+                return x
+
+        model = Net()
+
+        t = Variable(torch.ones(1, 3, 227, 227), requires_grad=True)
+
+        with torch.onnx.set_training(model, False):
+            trace, _ = torch.jit.trace(model, (t, ))
+
+        torch.onnx._optimize_trace(trace, False)
+        graph = torch._C._jit_get_graph(trace)
+        nodes = list(graph.nodes())
+
+        self.assertTrue(nodes[0].scopeName() == 'Net/Sequential[features]/Conv2d[0]')
+        self.assertTrue(nodes[1].scopeName() == 'Net/Sequential[features]/ReLU[1]')
+        self.assertTrue(nodes[2].scopeName() == 'Net/Sequential[features]/MaxPool2d[2]')
+
     @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows")
     @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA")
     def test_lstm_fusion(self):
