add tests and refactor code

lucent/modelzoo/util.py

@@ -16,15 +16,26 @@
 """Utility functions for modelzoo models."""
 
 from __future__ import absolute_import, division, print_function
+from collections import OrderedDict
 
 
-def get_model_layers(model):
-    layers = []
+def get_model_layers(model, getLayerRepr=False):
+    """
+    If getLayerRepr is True, return a OrderedDict of layer names, layer representation string pair.
+    If it's False, just return a list of layer names
+    """
+    layers = OrderedDict() if getLayerRepr else []
     # recursive function to get layers
     def get_layers(net, prefix=[]):
         if hasattr(net, "_modules"):
             for name, layer in net._modules.items():
-                layers.append("_".join(prefix+[name]))
+                if layer is None:
+                    # e.g. GoogLeNet's aux1 and aux2 layers
+                    continue
+                if getLayerRepr:
+                    layers["_".join(prefix+[name])] = layer.__repr__()
+                else:
+                    layers.append("_".join(prefix + [name]))
                 get_layers(layer, prefix=prefix+[name])
 
     get_layers(model)

lucent/util.py

@@ -19,7 +19,6 @@
 
 import torch
 import random
-from collections import OrderedDict
 
 
 def set_seed(seed):
@@ -28,18 +27,3 @@ def set_seed(seed):
     torch.cuda.manual_seed_all(seed)
     torch.backends.cudnn.deterministic=True
     random.seed(seed)
-
-def lucent_layernames(net, prefix=[]):
-    """ Return the layername and str representation of the layer """
-    layernames = OrderedDict()
-    def hook_layernames(net, prefix=[]):
-        """Recursive function to return the layer name"""
-        if hasattr(net, "_modules"):
-            for name, layer in net._modules.items():
-                if layer is None:
-                    # e.g. GoogLeNet's aux1 and aux2 layers
-                    continue
-                layernames["_".join(prefix+[name])] = layer.__repr__()
-                hook_layernames(layer, prefix=prefix+[name])
-    hook_layernames(net, prefix)
-    return layernames

tests/modelzoo/test_inceptionv1.py

@@ -39,3 +39,9 @@ def test_inceptionv1_graph_import():
     layer_names = util.get_model_layers(model)
     for layer_name in important_layer_names:
         assert layer_name in layer_names
+
+def test_inceptionv1_import_layer_repr():
+    model = inceptionv1()
+    layer_names = util.get_model_layers(model, getLayerRepr=True)
+    for layer_name in important_layer_names:
+        assert layer_names[layer_name] == 'CatLayer()'