Add main() for draw_net unittest, fix import errors

zhangrongen · Apr 4, 2017 · 31bfe8f · 31bfe8f
1 parent 41e34c9
commit 31bfe8f
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/python/caffe/test/test_draw.py b/python/caffe/test/test_draw.py
@@ -1,7 +1,7 @@
 import os
 import unittest
 
-from google import protobuf
+from google.protobuf import text_format
 
 import caffe.draw
 from caffe.proto import caffe_pb2
@@ -29,5 +29,9 @@ def test_draw_net(self):
         for filename in getFilenames():
             net = caffe_pb2.NetParameter()
             with open(filename) as infile:
-                protobuf.text_format.Merge(infile.read(), net)
+                text_format.Merge(infile.read(), net)
             caffe.draw.draw_net(net, 'LR')
+
+
+if __name__ == "__main__":
+    unittest.main()