Can take subview of container using []-syntax.

lucasb-eyer · Aug 31, 2015 · 15d0fbf · 15d0fbf
1 parent fcabd88
commit 15d0fbf
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/DeepFried2/containers/Container.py b/DeepFried2/containers/Container.py
@@ -1,5 +1,7 @@
 from ..layers import Module
 
+from DeepFried2.utils import aslist as _aslist
+
 
 class Container(Module):
 
@@ -38,3 +40,6 @@ def get_stat_updates(self):
     def add(self, *modules):
         assert all(isinstance(m, Module) for m in modules), "`Container`s can only contain objects subtyping `Module`."
         self.modules += modules
+
+    def __getitem__(self, slice_):
+        return type(self)(*_aslist(self.modules[slice_]))
diff --git a/DeepFried2/utils.py b/DeepFried2/utils.py
@@ -38,3 +38,12 @@ def load_params(module, fromwhere):
     with _np.load(fromwhere) as f:
         for p, v in zip(params, f['params']):
             p.set_value(v)
+
+
+def aslist(what):
+    if isinstance(what, list):
+        return what
+    elif isinstance(what, tuple):
+        return list(what)
+    else:
+        return [what]