Source code for megengine.module.sequential

# -*- coding: utf-8 -*-
from collections import OrderedDict

from .module import Module


[docs]class Sequential(Module): r"""A sequential container. Modules will be added to it in the order they are passed in the constructor. Alternatively, an ordered dict of modules can also be passed in. Examples: .. testcode:: import numpy as np import megengine as mge import megengine.module as M import megengine.functional as F from collections import OrderedDict batch_size = 64 data = mge.tensor(np.zeros((batch_size, 28 * 28)), dtype=np.float32) label = mge.tensor(np.zeros(batch_size,), dtype=np.int32) net0 = M.Sequential( M.Linear(28 * 28, 320), M.Linear(320, 10) ) pred0 = net0(data) modules = OrderedDict() modules["fc0"] = M.Linear(28 * 28, 320) modules["fc1"] = M.Linear(320, 10) net1 = M.Sequential(modules) pred1 = net1(data) """ def __init__(self, *args, **kwargs): super().__init__(**kwargs) self.layer_keys = [] if len(args) == 1 and isinstance(args[0], OrderedDict): for key, module in args[0].items(): # self.add_module(key, module) setattr(self, key, module) self.layer_keys.append(key) else: for idx, module in enumerate(args): # self.add_module(str(idx), module) setattr(self, str(idx), module) self.layer_keys.append(str(idx)) def __getitem__(self, idx): if isinstance(idx, slice): return self.__class__( OrderedDict(zip(self.layer_keys[idx], self.layer_values[idx])) ) else: return getattr(self, self.layer_keys[idx]) def __setitem__(self, idx, module): key = self.layer_keys[idx] return setattr(self, key, module) def __delitem__(self, idx): if isinstance(idx, slice): for key in self.layer_keys[idx]: delattr(self, key) del self.layer_keys[idx] else: delattr(self, self.layer_keys[idx]) del self.layer_keys[idx] def __len__(self): return len(self.layer_keys) def __iter__(self): return iter(self.layer_values) @property def layer_values(self): return [getattr(self, key) for key in self.layer_keys] def forward(self, inp): # avoid layer_values as a name prefix, see Module.__getattribute__ for layer in [getattr(self, key) for key in self.layer_keys]: inp = layer(inp) return inp