Метод add_module

Метод add_module используется для динамического добавления дочернего модуля в текущий модуль. Первым параметром метод принимает имя подмодуля в виде строки, вторым параметром - добавляемый модуль. Это позволяет гибко конструировать архитектуру нейронных сетей во время выполнения программы.

Синтаксис

module.add_module(name, submodule)

Пример

Давайте добавим линейный слой в пустой модуль:

import torch
import torch.nn as nn

model = nn.Module()
model.add_module('fc1', nn.Linear(10, 20))
print(model)

Результат выполнения кода:

Module(
  (fc1): Linear(in_features=10, out_features=20, bias=True)
)

Пример

Давайте последовательно добавим несколько слоёв в модель:

import torch
import torch.nn as nn

model = nn.Module()
model.add_module('fc1', nn.Linear(10, 20))
model.add_module('fc2', nn.Linear(20, 30))
model.add_module('relu', nn.ReLU())
print(model)

Результат выполнения кода:

Module(
  (fc1): Linear(in_features=10, out_features=20, bias=True)
  (fc2): Linear(in_features=20, out_features=30, bias=True)
  (relu): ReLU()
)

Пример

Давайте динамически создадим последовательную сеть с помощью цикла:

import torch
import torch.nn as nn

model = nn.Module()
layer_sizes = [10, 20, 30, 10]

for i in range(len(layer_sizes) - 1):
	layer_name = f'linear_{i}'
	model.add_module(
		layer_name,
		nn.Linear(layer_sizes[i], layer_sizes[i + 1])
	)
	model.add_module(f'relu_{i}', nn.ReLU())

print(model)

Результат выполнения кода:

Module(
  (linear_0): Linear(in_features=10, out_features=20, bias=True)
  (relu_0): ReLU()
  (linear_1): Linear(in_features=20, out_features=30, bias=True)
  (relu_1): ReLU()
  (linear_2): Linear(in_features=30, out_features=10, bias=True)
  (relu_2): ReLU()
)

Пример

Давайте создадим модуль с вложенным модулем, используя add_module:

import torch
import torch.nn as nn

class MyModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.add_module('subnet', nn.Sequential(
			nn.Linear(10, 20),
			nn.ReLU(),
			nn.Linear(20, 5)
		))
		self.add_module('final', nn.Linear(5, 1))

model = MyModel()
print(model)

Результат выполнения кода:

MyModel(
  (subnet): Sequential(
	(0): Linear(in_features=10, out_features=20, bias=True)
	(1): ReLU()
	(2): Linear(in_features=20, out_features=5, bias=True)
  )
  (final): Linear(in_features=5, out_features=1, bias=True)
)

Смотрите также

метод register_module,
который выполняет ту же функцию, что и add_module
метод children,
который возвращает итератор по дочерним модулям
метод named_children,
который возвращает итератор по именам и дочерним модулям
метод modules,
который возвращает итератор по всем модулям в иерархии