Как Pytorch's nn.Module регистрирует подмодуль?

Question

Когда я читал исходный код (python) файла torch.nn.Module, я обнаружил, что атрибут self._modules используется во многих функциях, таких как self.modules(), self.children() и т. Д. Однако я не нашел ни одногофункции, обновляющие его.Итак, где будет обновляться self._modules?Кроме того, как Pytorch nn.Module регистрирует субмодуль?

class Module(object):
    def __init__(self):
        self._backend = thnn_backend
        self._parameters = OrderedDict()
        self._buffers = OrderedDict()
        self._backward_hooks = OrderedDict()
        self._forward_hooks = OrderedDict()
        self._forward_pre_hooks = OrderedDict()
        self._modules = OrderedDict()
        self.training = True

    def named_modules(self, memo=None, prefix=''):
        if memo is None:
            memo = set()
        if self not in memo:
            memo.add(self)
            yield prefix, self
            for name, module in self._modules.items():
                if module is None:
                    continue
                submodule_prefix = prefix + ('.' if prefix else '') + name
                for m in module.named_modules(memo, submodule_prefix):
                    yield m

Answer 1

Добавьте некоторые детали к ответу Джирена Джина:

• Слои сети (унаследованные от nn.Module) хранятся в Module._modules, который инициализируется в __construct:

  def __init__(self):
      self.__construct()
      # initialize self.training separately from the rest of the internal
      # state, as it is managed differently by nn.Module and ScriptModule
      self.training = True
  
  def __construct(self):
      """
      Initializes internal Module state, shared by both nn.Module and ScriptModule.
      """
      # ...
      self._modules = OrderedDict()
  

• self._modules обновлено в __setattr__.__setattr__(obj, name, value) вызывается при выполнении obj.name = value.Например, если при инициализации сети, унаследованной от nn.Module, определяется self.conv1 = nn.Conv2d(128, 256, 3, 1, 1), будет выполнен следующий код из nn.Module.__setattr__:

  def __setattr__(self, name, value):
      def remove_from(*dicts):
          for d in dicts:
              if name in d:
                  del d[name]
  
      params = self.__dict__.get('_parameters')
      if isinstance(value, Parameter):
          # ...
      elif params is not None and name in params:
          # ...
      else:
          modules = self.__dict__.get('_modules') # equivalent to modules = self._modules
          if isinstance(value, Module):
              if modules is None:
                  raise AttributeError(
                      "cannot assign module before Module.__init__() call")
              remove_from(self.__dict__, self._parameters, self._buffers)
              # register the given layer (nn.Conv2d) with its name (conv1)
              # equivalent to self._modules['conv1'] = nn.Conv2d(128, 256, 3, 1, 1)
              modules[name] = value
  

Вопрос из комментариев:

Знаете ли вы, как это работает с тем фактом, что torch позволяет вам использовать свой собственный метод пересылки?

Если выполняется прямой пересылка сети, унаследованной от nn.Moduleбудет вызываться nn.Module.__call__, в котором вызывается self.forward.Тем не менее, один переопределил forward при реализации сети.

Answer 2

Модули и параметры обычно регистрируются путем установки атрибута для экземпляра nn.module.В частности, этот тип поведения реализуется путем использования метода __serattr__:

def __setattr__(self, name, value):
        def remove_from(*dicts):
            for d in dicts:
                if name in d:
                    del d[name]

        params = self.__dict__.get('_parameters')
        if isinstance(value, Parameter):
            if params is None:
                raise AttributeError(
                    "cannot assign parameters before Module.__init__() call")
            remove_from(self.__dict__, self._buffers, self._modules)
            self.register_parameter(name, value)
        elif params is not None and name in params:
            if value is not None:
                raise TypeError("cannot assign '{}' as parameter '{}' "
                                "(torch.nn.Parameter or None expected)"
                                .format(torch.typename(value), name))
            self.register_parameter(name, value)
        else:
            modules = self.__dict__.get('_modules')
            if isinstance(value, Module):
                if modules is None:
                    raise AttributeError(
                        "cannot assign module before Module.__init__() call")
                remove_from(self.__dict__, self._parameters, self._buffers)
                modules[name] = value
            elif modules is not None and name in modules:
                if value is not None:
                    raise TypeError("cannot assign '{}' as child module '{}' "
                                    "(torch.nn.Module or None expected)"
                                    .format(torch.typename(value), name))
                modules[name] = value
            else:
                buffers = self.__dict__.get('_buffers')
                if buffers is not None and name in buffers:
                    if value is not None and not isinstance(value, torch.Tensor):
                        raise TypeError("cannot assign '{}' as buffer '{}' "
                                        "(torch.Tensor or None expected)"
                                        .format(torch.typename(value), name))
                    buffers[name] = value
                else:
                    object.__setattr__(self, name, value)

См. https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/module.py, чтобы найти этот метод.