Feature request: ModuleDict, like ModuleList

[rkaplan]

Self-explanatory. Currently there is no easy way to maintain a dictionary as an attribute of a Module whose values are themselves Modules that need to be registered. Introducing a ModuleDict class that functions the same as a ModuleList, except it exposes a dictionary interface, would resolve this. This has the following practical benefits for users:

1. Checkpointing is much better with a ModuleDict that has strings as keys and Modules as values than a ModuleList. I have personally run into problems when I change the order of modules in my ModuleList and try to load old checkpoints, because the checkpoint key names are order-dependent with ModuleLists, so it breaks my code.

2. It combines the advantages of a ModuleList (dynamic assignment of Modules to a class) with the advantages of using named attributes (documentation / readability / sane module names in checkpoints).

Another pro is that this is very simple to implement. Here is a first pass, if there is support for this feature I will clean it up and submit this as a PR:

class ModuleDict(Module):
    r"""Holds submodules in a dict.
    ModuleDict can be indexed like a regular Python dict, but modules it
    contains are properly registered, and will be visible by all Module methods.
    Arguments:
        modules (dict, optional): a list of modules to add
    Example::
        TODO.
    """

    def __init__(self, modules=None):
        super(ModuleDict, self).__init__()
        if modules is not None:
            self.update(modules)

    def __getitem__(self, key):
        return self._modules[key]

    def __setitem__(self, key, module):
        return setattr(self, key, module)

    def __len__(self):
        return len(self._modules)

    def __iter__(self):
        return iter(self._modules)

    def keys(self):
        return self._modules.keys()

    def items(self):
        return self._modules.items()

    def values(self):
        return self._modules.values()

    def get(self, key, default=None):
        if default is not None:
            return self._modules.get(key, default)
        return self._modules.get(key)

    def update(self, modules):
        r"""Updates modules from a Python dict.
        Arguments:
            modules (dict): dict of modules to append
        """
        if not isinstance(modules, dict):
            raise TypeError("ModuleDict.update should be called with a "
                            "dict, but got " + type(modules).__name__)
        for key, module in modules.items():
            self.add_module(key, module)
        return self

[dhpollack]

You should also allow the named_module() generator as an input to the update function. Also would be good to have this with parameters as well.

[apaszke]

Thanks for the proposal. I’m wondering if there’s any reason why you can just use nn.Module and regular attribute access for this purpose?

[dhpollack]

@apaszke, I think you are right to say that the normal nn.Module should be used for this, but based on the documentation it is not obvious that one can use the nn.Module for this purpose. I think if there was an example in the nn.ModuleList and/or nn.Module documentation that showed how nn.Module could be used as a "ModuleList dictionary", then one wouldn't need or request a nn.ModuleDict.

[fmassa]

Here is another proposal:

Why not create a nn.ModuleIterable (or something like that) that wraps an iterable into a nn.Module, like what is currently done in nn.ModuleList, and deprecate nn.ModuleList?

I'm not sure if it would be best to add a number of nn.Modules for different container types.

What do you think?

[apaszke]

I’m not sure how does this solve the problem or how is it different than ModuleList? Can you please elaborate on your idea?

[fmassa]

Implementation-wise this would be very similar to what we currently have for nn.ModuleList, eventually bookeeping the type of the iterable that was passed so that we can provide nice printings and indexing access.

This solves the problem because every iterable that you want to pass you just need to pass that you want to record in the model (for .parameters(), .cuda(), etc) to this common module.

Is such a thing necessary? Not necessarily, but it was already handy in the past.

Do we need such abstraction? Maybe not, as we can subclass a nn.Module to implement such funcionality. But for users, might be better to do

branch1 = {}
branch1['conv1'] = nn.Conv2d(...)
branch1['conv2'] = nn.Conv2d(...)
branch1['conv3'] = nn.Conv2d(...)
self.branch1 = nn.FromIterable(branch1)

# and then later on
for module_name in ['conv1', 'conv2', 'conv3']:
# or for module in self.branch1 if we branch1 was OrderedDict
    output = F.relu(self.branch1[module_name]) + 1.0

instead of having to write a class just for holding the conv1/conv2/conv3.

Let me know if I wasn't clear.

[rkaplan]

I prefer a dict interface to using regular attribute access with nn.Module because in my case, the names of the modules are only known at runtime. (They are dependent on the config file). I have a multi-headed network with several heads doing different tasks; each head is a Module that I'd like to keep as the value of a dictionary whose key is the task name.

Furthermore, due to the way my data is labeled, I can only backprop through one head at a time. For each image I only have one of the tasks labeled. So each minibatch requires forwarding through a different subset of the network. This is a primary reason I'm using PyTorch, and I think multi-task learning setups like this will only get more common. To choose which subset of the network the data gets forwarded through, I use a string provided by the data loader: ideally the string would just be the key for the heads ModuleDict.

It's possible to do this using getattr and setattr, but given the amount of magic PyTorch does for these functions in nn.Module under the hood, it's preferable to have a simple dict abstraction. It seems like one of PyTorch's design goals is first-class support for dynamic graphs -- if that's true then a ModuleDict is a natural addition.

[dhpollack]

I agree it is not necessary, but does seem like a "nice to have" addition.

Another point, is that nn.ModuleDict is more natural than nn.ModuleList, because modules and parameters are stored in an OrderedDict already.

[apaszke]

@fmassa I guess my confusion comes from the fact that Iterable in Python means just a (possibly lazy) sequence of items, with specified order and without the mapping structure. This is exactly what ModuleList provides.

I wouldn't say that it's more natural - both data structures have their own different purposes, but after giving it some thought I think I'm ok with adding ModuleDict, so please send a PR.

[fmassa]

@apaszke maybe my confusion comes from the fact that

isinstance({}, collections.Iterable) == True

Anyway, my idea was just to unite all those new modules in a single one, so that we don't grow it for each new container type (ModuleNamedTuple, ModuleSet, etc).

But that's just a minor comment.

[apaszke]

Well that's because iterable is just something you can iterate over, but e.g. iterating over a dict only gives you keys by default (iterable doesn't imply that you need to recover the full container structure). In general, every available container is different in some way, and we would probably want to keep these differences when enabling their use as containers for modules (e.g. insert(n, item) is only valid for lists, keys is only valid for dicts). I don't think there's a good way to unify them.

[apaszke]

Also, lists and dicts are kind of fundamental. We won't be getting a named tuple or a set for sure (because comparison on modules is not very well defined, and you'd be better off using a dict with meaningful keys).

[ozancaglayan]

I have a use case very similar to what has been described by @rkaplan . I may have N encoders and M decoders depending on the configuration file which defines several datasets let's say. I would like to select a specific encoder and/or decoder based on a string that'll change during training.

The easiest way I'd imagine for doing this is sth like below in the init of the custom Module:

self.encoders = {}
for enc in requested_encoders:
  setattr(self, 'encoder_%s' % enc, RNN(..))
  self.encoders[enc] = getattr(self, 'encoder_%s' % enc)

# Or alternatively if the Module could (semi-)automatically detect dicts
self.encoders = {enc:RNN(..) for enc in requested_encoders}

[apaszke]

@ozancaglayan this will break if you try to use it with DataParallel. We should just get a ModuleDict, but as a workaround, you can keep a module list, and make the dict map names to indices in that list (this will work with DataParallel)

[ozancaglayan]

Just out of curiosity, why does this break DataParallel?