Список слоёв в PyTorch

Для цепочки однотипных шагов удобен класс ModuleList: линейные слои лежат в списке, а в forward вызывают нужный элемент по индексу. Параметры таких слоёв попадают в общий перебор модуля.

Соберём две линейные связи в списке и прогоним через них вектор входа:

import torch
import torch.nn as nn

class StackNet(nn.Module):
	def __init__(self):
		super().__init__()
		self.layers = nn.ModuleList([
			nn.Linear(2, 4),
			nn.Linear(4, 2),
		])

	def forward(self, x):
		x = self.layers[0](x)
		return self.layers[1](x)

net = StackNet()
x = torch.tensor([[1.0, 2.0]])
out = net(x)
print(out.shape) # выведет torch.Size([1, 2])

Число обучаемых элементов у такой модели включает веса обоих слоёв из списка:

import torch
import torch.nn as nn

class StackNet(nn.Module):
	def __init__(self):
		super().__init__()
		self.layers = nn.ModuleList([
			nn.Linear(2, 4),
			nn.Linear(4, 2),
		])

	def forward(self, x):
		x = self.layers[0](x)
		return self.layers[1](x)

net = StackNet()
total = sum(p.numel() for p in net.parameters())
print(total) # выведет 22

Опишите модуль с двумя линейными слоями 3 на 2 и 2 на 1 в списке. Вызовите элементы по индексу в прямом проходе и выведите форму ответа на вход [[1.0, 2.0, 0.5]].

Создайте блок с тремя линейными преобразованиями 2 на 2 в списке из трёх элементов. Выведите общее число обучаемых скаляров у всего модуля.

Соберите сеть: список из линейного слоя 1 на 4 и второго 4 на 2. Выведите, сколько отдельных тензоров параметров видит корневой модуль.