Класс Module

Класс Module является базовым классом для всех нейросетевых слоёв в PyTorch. Любой пользовательский слой или модель должны наследоваться от этого класса. Основная задача класса Module - предоставить единый интерфейс для построения нейросетей, управления параметрами, переключения режимов обучения и сохранения состояния. При создании собственного слоя необходимо переопределить метод forward, в котором описывается преобразование данных.

Синтаксис

class MyLayer(torch.nn.Module):
	def __init__(self):
		super().__init__()
		# инициализация параметров и подслоёв
	
	def forward(self, x):
		# преобразование входных данных
		return x

Пример

Давайте создадим простой полносвязный слой без смещения:

import torch

class LinearLayer(torch.nn.Module):
	def __init__(self, in_features, out_features):
		super().__init__()
		self.weight = torch.nn.Parameter(
			torch.randn(in_features, out_features)
		)
	
	def forward(self, x):
		return x @ self.weight

layer = LinearLayer(3, 2)
t = torch.tensor([1.0, 2.0, 3.0])
res = layer(t)
print(res)

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

tensor([-0.8748, -3.0315], grad_fn=<SqueezeBackward3>)

Пример

Создадим простую нейросеть из двух слоёв с функцией активации:

import torch

class SimpleNet(torch.nn.Module):
	def __init__(self):
		super().__init__()
		self.fc1 = torch.nn.Linear(10, 5)
		self.fc2 = torch.nn.Linear(5, 1)
		self.relu = torch.nn.ReLU()
	
	def forward(self, x):
		x = self.relu(self.fc1(x))
		x = self.fc2(x)
		return x

model = SimpleNet()
t = torch.randn(3, 10)
res = model(t)
print(res.shape)

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

torch.Size([3, 1])

Пример

Используем метод train для переключения в режим обучения:

import torch

class DropoutNet(torch.nn.Module):
	def __init__(self):
		super().__init__()
		self.fc = torch.nn.Linear(5, 5)
		self.dropout = torch.nn.Dropout(0.5)
	
	def forward(self, x):
		x = self.fc(x)
		x = self.dropout(x)
		return x

model = DropoutNet()
print(f"train mode: {model.training}")
model.train()
print(f"train mode: {model.training}")

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

"train mode: True"
"train mode: True"

Пример

Переключим модель в режим оценки с помощью метода eval:

import torch

model = DropoutNet()
print(f"train mode: {model.training}")
model.eval()
print(f"train mode: {model.training}")

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

"train mode: True"
"train mode: False"

Пример

Получим все параметры модели с помощью метода parameters:

import torch

model = SimpleNet()
for name, param in model.named_parameters():
	print(f"{name}: {param.shape}")

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

"fc1.weight: torch.Size([5, 10])"
"fc1.bias: torch.Size([5])"
"fc2.weight: torch.Size([1, 5])"
"fc2.bias: torch.Size([1])"

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

метод forward,
который определяет преобразование данных в слое
метод train,
который переключает модуль в режим обучения
метод eval,
который переключает модуль в режим оценки
метод parameters,
который возвращает итератор по параметрам модуля