make layer_names dynamically changeable

chainercv/links/model/sequential_feature_extractor.py

@@ -1,5 +1,4 @@
 import collections
-from itertools import islice
 
 import chainer
 
@@ -38,16 +37,13 @@ class SequentialFeatureExtractor(chainer.Chain):
         >>>     ('l2', L.Linear(None, 1000)),
         >>>     ('l2_relu', F.relu),
         >>>     ('l3', L.Linear(None, 10))])
-        >>> model = SequentialFeatureExtractor(layers, ['l2_relu', 'l3'])
-        >>> # These are outputs of l2_relu and l3 layers.
-        >>> feat1, feat2 = model(imgs)
-
-
-    The implementation is optimized for speed and memory.
-    A layer that is not needed to collect all features listed in
-    :obj:`layer_names` will not be added as a child link.
-    Also, this object only conducts the minimal amount of computation needed
-    to collect these features.
+        >>> model = SequentialFeatureExtractor(layers, ['l2_relu', 'l1_relu'])
+        >>> # These are outputs of layer l2_relu and l1_relu.
+        >>> feat1, feat2 = model(x)
+        >>> # The layer_names can be dynamically changed.
+        >>> model.layer_names = 'l3'
+        >>> # This is an output of layer l1.
+        >>> feat3 = model(x)
 
     Args:
         layers (list or collections.OrderedDict of callables):
@@ -67,6 +63,19 @@ def __init__(self, layers, layer_names=None):
                  for i, layer in enumerate(layers)])
         self._layers = layers
 
+        self.layer_names = layer_names
+
+        with self.init_scope():
+            for name, layer in self._layers.items():
+                if isinstance(layer, chainer.Link):
+                    setattr(self, name, layer)
+
+    @property
+    def layer_names(self):
+        return self._layer_names
+
+    @layer_names.setter
+    def layer_names(self, layer_names):
         if layer_names is None:
             layer_names = list(self._layers.keys())[-1]
 
@@ -76,21 +85,12 @@ def __init__(self, layers, layer_names=None):
         else:
             return_tuple = False
             layer_names = [layer_names]
+        if any([name not in self._layers for name in layer_names]):
+            raise ValueError('Invalid layer name')
+
         self._return_tuple = return_tuple
         self._layer_names = layer_names
 
-        # Delete unnecessary layers from self._layers based on layer_names.
-        # Computation is equivalent to layers = layers[:last_index + 1].
-        last_index = max([list(self._layers.keys()).index(name) for
-                          name in self._layer_names])
-        self._layers = collections.OrderedDict(
-            islice(self._layers.items(), None, last_index + 1))
-
-        with self.init_scope():
-            for name, layer in self._layers.items():
-                if isinstance(layer, chainer.Link):
-                    setattr(self, name, layer)
-
     def __call__(self, x):
         """Forward sequential feature extraction model.
 
@@ -102,9 +102,14 @@ def __call__(self, x):
             The returned values are determined by :obj:`layer_names`.
 
         """
+        # The biggest index among indices of the layers that are included
+        # in self._layer_names.
+        last_index = max([list(self._layers.keys()).index(name) for
+                          name in self._layer_names])
+
         features = {}
         h = x
-        for name, layer in self._layers.items():
+        for name, layer in list(self._layers.items())[:last_index + 1]:
             h = layer(h)
             if name in self._layer_names:
                 features[name] = h

tests/links_tests/model_tests/test_sequential_feature_extractor.py

@@ -34,30 +34,26 @@ def setUp(self):
                  ('f1', self.f1),
                  ('f2', self.f2),
                  ('l2', self.l2)]),
-            layer_names=['l1', 'f1', 'f2', 'l2'])
+            layer_names=['l1', 'f1', 'f2'])
         self.x = np.random.uniform(size=(1, 3, 24, 24))
 
     def check_call_output(self):
         x = self.link.xp.asarray(self.x)
         out = self.link(x)
 
-        self.assertEqual(len(out), 4)
+        self.assertEqual(len(out), 3)
         self.assertIsInstance(out[0], chainer.Variable)
         self.assertIsInstance(out[1], chainer.Variable)
         self.assertIsInstance(out[2], chainer.Variable)
-        self.assertIsInstance(out[3], chainer.Variable)
         self.assertIsInstance(out[0].data, self.link.xp.ndarray)
         self.assertIsInstance(out[1].data, self.link.xp.ndarray)
         self.assertIsInstance(out[2].data, self.link.xp.ndarray)
-        self.assertIsInstance(out[3].data, self.link.xp.ndarray)
 
         out_data = [to_cpu(var.data) for var in out]
         np.testing.assert_equal(out_data[0], to_cpu(self.l1(x).data))
         np.testing.assert_equal(out_data[1], to_cpu(self.f1(self.l1(x)).data))
         np.testing.assert_equal(
             out_data[2], to_cpu(self.f2(self.f1(self.l1(x))).data))
-        np.testing.assert_equal(
-            out_data[3], to_cpu(self.l2(self.f2(self.f1(self.l1(x)))).data))
 
     def test_call_output_cpu(self):
         self.check_call_output()
@@ -67,6 +63,25 @@ def test_call_output_gpu(self):
         self.link.to_gpu()
         self.check_call_output()
 
+    def check_call_dynamic_layer_names(self):
+        x = self.link.xp.asarray(self.x)
+        self.link.layer_names = ['l2']
+        out, = self.link(x)
+
+        self.assertIsInstance(out, chainer.Variable)
+        self.assertIsInstance(out.data, self.link.xp.ndarray)
+
+        out_data = out.data
+        np.testing.assert_equal(
+            out_data, to_cpu(self.l2(self.f2(self.f1(self.l1(x)))).data))
+
+    def test_call_dynamic_layer_names_cpu(self):
+        self.check_call_dynamic_layer_names()
+
+    @attr.gpu
+    def test_call_dynamic_layer_names_gpu(self):
+        self.check_call_dynamic_layer_names()
+
 
 class TestSequentialFeatureExtractorListFunctions(unittest.TestCase):
 
@@ -131,32 +146,4 @@ def test_copy_gpu(self):
         self.check_copy()
 
 
-class TestSequentialFeatureExtractorRedundantLayers(unittest.TestCase):
-
-    def setUp(self):
-        self.l1 = ConstantStubLink(np.random.uniform(size=(1, 3, 24, 24)))
-        self.f1 = DummyFunc()
-        self.f2 = DummyFunc()
-        self.l2 = ConstantStubLink(np.random.uniform(size=(1, 3, 24, 24)))
-
-        self.link = SequentialFeatureExtractor(
-            collections.OrderedDict(
-                [('l1', self.l1),
-                 ('f1', self.f1),
-                 ('f2', self.f2),
-                 ('l2', self.l2)]),
-            layer_names=['l1', 'f1'])
-
-    def check_redundant_layers(self):
-        self.assertNotIn('f2', self.link._layer_names)
-        self.assertNotIn('l2', self.link._layer_names)
-
-    def test_redundant_layers_cpu(self):
-        self.check_redundant_layers()
-
-    @attr.gpu
-    def test_redundant_layers_gpu(self):
-        self.check_redundant_layers()
-
-
 testing.run_module(__name__, __file__)