add sequential_feature_extractor

chainercv/links/__init__.py

@@ -1,4 +1,5 @@
 from chainercv.links.model.pixelwise_softmax_classifier import PixelwiseSoftmaxClassifier  # NOQA
+from chainercv.links.model.sequential_feature_extractor import SequentialFeatureExtractor  # NOQA
 
 from chainercv.links.model.faster_rcnn.faster_rcnn_vgg import FasterRCNNVGG16  # NOQA
 from chainercv.links.model.segnet.segnet_basic import SegNetBasic  # NOQA

chainercv/links/model/sequential_feature_extractor.py

@@ -0,0 +1,132 @@
+import collections
+
+import chainer
+
+
+class SequentialFeatureExtractor(chainer.Chain):
+
+    """A feature extractor model with a single-stream forward pass.
+
+    This class is a base class that can be used for an implementation of
+    a feature extractor model.
+    The link takes :obj:`layers` to specify the computation
+    conducted in :meth:`__call__`.
+    :obj:`layers` is a list or :obj:`collections.OrderedDict` of
+    callable objects called layers, which are going to be called sequentially
+    starting from the top to the end.
+    A :obj:`chainer.Link` object in the sequence will be added as
+    a child link of this object.
+
+    :meth:`__call__` returns single or multiple features that are picked up
+    through a stream of computation.
+    These features can be specified by :obj:`layer_names`, which contains
+    the names of the layer whose output is collected.
+    When :obj:`layer_names` is a string, single value is returned.
+    When :obj:`layer_names` is an iterable of strings, a tuple of values
+    will be returned. These values are ordered in the same order of the
+    strings in :obj:`layer_names`.
+
+    Examples:
+
+        >>> import collections
+        >>> import chainer.functions as F
+        >>> import chainer.links as L
+        >>> layers = collections.OrderedDict([
+        >>>     ('l1', L.Linear(None, 1000)),
+        >>>     ('l1_relu', F.relu),
+        >>>     ('l2', L.Linear(None, 1000)),
+        >>>     ('l2_relu', F.relu),
+        >>>     ('l3', L.Linear(None, 10))])
+        >>> 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):
+            Callable objects called in the forward pass.
+        layer_names (string or iterable of strings):
+            Names of layers whose outputs will be collected in
+            the forward pass.
+
+    """
+
+    def __init__(self, layers, layer_names=None):
+        super(SequentialFeatureExtractor, self).__init__()
+
+        if not isinstance(layers, collections.OrderedDict):
+            layers = collections.OrderedDict(
+                [('{}_{}'.format(layer.__class__.__name__, i), layer)
+                 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]
+
+        if (not isinstance(layer_names, str) and
+                all([isinstance(name, str) for name in layer_names])):
+            return_tuple = True
+        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
+
+    def __call__(self, x):
+        """Forward sequential feature extraction model.
+
+        Args:
+            x (chainer.Variable or array): Input to the network.
+
+        Returns:
+            chainer.Variable or tuple of chainer.Variable:
+            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 list(self._layers.items())[:last_index + 1]:
+            h = layer(h)
+            if name in self._layer_names:
+                features[name] = h
+
+        if self._return_tuple:
+            features = tuple(
+                [features[name] for name in self._layer_names])
+        else:
+            features = list(features.values())[0]
+        return features
+
+    def copy(self):
+        ret = super(SequentialFeatureExtractor, self).copy()
+        layers = []
+        for name, layer in self._layers.items():
+            if name in self._children:
+                layer = ret[name]
+            layers.append((name, layer))
+        ret.layers = collections.OrderedDict(layers)
+        return ret

tests/links_tests/model_tests/test_sequential_feature_extractor.py

@@ -0,0 +1,149 @@
+import collections
+import unittest
+
+import numpy as np
+
+import chainer
+from chainer.cuda import to_cpu
+from chainer import testing
+from chainer.testing import attr
+
+from chainer.function import Function
+
+from chainercv.links import SequentialFeatureExtractor
+from chainercv.utils.testing import ConstantStubLink
+
+
+class DummyFunc(Function):
+
+    def forward(self, inputs):
+        return inputs[0] * 2,
+
+
+class TestSequentialFeatureExtractorOrderedDictFunctions(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', '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), 3)
+        self.assertIsInstance(out[0], chainer.Variable)
+        self.assertIsInstance(out[1], chainer.Variable)
+        self.assertIsInstance(out[2], 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)
+
+        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))
+
+    def test_call_output_cpu(self):
+        self.check_call_output()
+
+    @attr.gpu
+    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):
+
+    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(
+            [self.l1, self.f1, self.f2, self.l2])
+        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.assertIsInstance(out, chainer.Variable)
+        self.assertIsInstance(out.data, self.link.xp.ndarray)
+
+        out = to_cpu(out.data)
+        np.testing.assert_equal(
+            out,
+            to_cpu(self.l2(self.f2(self.f1(self.l1(x)))).data))
+
+    def test_call_output_cpu(self):
+        self.check_call_output()
+
+    @attr.gpu
+    def test_call_output_gpu(self):
+        self.link.to_gpu()
+        self.check_call_output()
+
+
+class TestSequentialFeatureExtractorCopy(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', 'f2', 'l2'])
+
+    def check_copy(self):
+        copied = self.link.copy()
+        self.assertIs(copied.l1, copied.layers['l1'])
+        self.assertIs(copied.l2, copied.layers['l2'])
+
+    def test_copy_cpu(self):
+        self.check_copy()
+
+    @attr.gpu
+    def test_copy_gpu(self):
+        self.link.to_gpu()
+        self.check_copy()
+
+
+testing.run_module(__name__, __file__)