Функция torch.nn.Module в PyTorch
Справочное руководство PyTorch torch.nn
torch.nn.ModuleЭто базовый класс для всех модулей нейронных сетей в PyTorch.
Все пользовательские модели сетей должны наследоваться от этого класса, так как он предоставляет управление параметрами, перенос на устройство, сохранение модели и другие функции.
Определение класса
torch.nn.Module
Основные атрибуты
parameters(): Возвращает все обучаемые параметры модели.named_parameters(): Возвращает итератор по именам и значениям параметров.children(): Возвращает подмодули модели.named_children(): Возвращает итератор, содержащий имена подмодулей и модули.modules(): Возвращает все модули.state_dict(): Возвращает словарь, содержащий все параметры.
Примеры использования
Пример 1: Создание пользовательского модуля
Наследуйте nn.Module, чтобы создать пользовательскую сеть:
Пример
import torch
import torch.nn as nn
class SimpleNet(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super(SimpleNet, self).__init__()
# Определить слои сети
self.fc1 = nn.Linear(input_dim, hidden_dim)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_dim, output_dim)
def forward(self, x):
# Определить прямое распространение
x = self.fc1(x)
x = self.relu(x)
x = self.fc2(x)
return x
# Создать экземпляр модели
model = SimpleNet(input_dim=784, hidden_dim=256, output_dim=10)
#
x = torch.randn(32, 784)
output = model(x)
print(Структура модели:)
print(model)
print(nФорма входных данных:, x.shape)
print(Форма выходных данных:, output.shape)
import torch.nn as nn
class SimpleNet(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super(SimpleNet, self).__init__()
# Определить слои сети
self.fc1 = nn.Linear(input_dim, hidden_dim)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_dim, output_dim)
def forward(self, x):
# Определить прямое распространение
x = self.fc1(x)
x = self.relu(x)
x = self.fc2(x)
return x
# Создать экземпляр модели
model = SimpleNet(input_dim=784, hidden_dim=256, output_dim=10)
#
x = torch.randn(32, 784)
output = model(x)
print(Структура модели:)
print(model)
print(nФорма входных данных:, x.shape)
print(Форма выходных данных:, output.shape)
Пример 2: Управление параметрами
Доступ и управление параметрами модели:
Пример
import torch
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 16, 3)
self.bn1 = nn.BatchNorm2d(16)
self.fc = nn.Linear(16 * 14 * 14, 10)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
model = Net()
# Подсчет количества параметров
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(Всего параметров:, total_params)
print(Обучаемые параметры:, trainable_params)
# Доступ к конкретным параметрам
print(nconv1 форма веса:, model.conv1.weight.shape)
print(nfc форма смещения:, model.fc.bias.shape)
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 16, 3)
self.bn1 = nn.BatchNorm2d(16)
self.fc = nn.Linear(16 * 14 * 14, 10)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
model = Net()
# Подсчет количества параметров
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(Всего параметров:, total_params)
print(Обучаемые параметры:, trainable_params)
# Доступ к конкретным параметрам
print(nconv1 форма веса:, model.conv1.weight.shape)
print(nfc форма смещения:, model.fc.bias.shape)
Пример 3: Сохранение и загрузка модели
Сохраните и загрузите модель:
Пример
import torch
import torch.nn as nn
import tempfile
import os
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc = nn.Linear(10, 5)
def forward(self, x):
return self.fc(x)
model = Net()
# Сохранить всю модель
with tempfile.NamedTemporaryFile(delete=False, suffix='.pth') as f:
torch.save(model, f.name)
path_full = f.name
# Сохранить state_dict (рекомендуемый метод)
with tempfile.NamedTemporaryFile(delete=False, suffix='.pth') as f:
torch.save(model.state_dict(), f.name)
path_state = f.name
# Загрузить модель
loaded_model = Net()
loaded_model.load_state_dict(torch.load(path_state))
loaded_model.eval()
# Тестирование загруженной модели
x = torch.randn(2, 10)
output1 = model(x)
output2 = loaded_model(x)
print(Исходный вывод:, output1[0].tolist())
print(Вывод после загрузки:, output2[0].tolist())
# Очистка
os.remove(path_full)
os.remove(path_state)
import torch.nn as nn
import tempfile
import os
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc = nn.Linear(10, 5)
def forward(self, x):
return self.fc(x)
model = Net()
# Сохранить всю модель
with tempfile.NamedTemporaryFile(delete=False, suffix='.pth') as f:
torch.save(model, f.name)
path_full = f.name
# Сохранить state_dict (рекомендуемый метод)
with tempfile.NamedTemporaryFile(delete=False, suffix='.pth') as f:
torch.save(model.state_dict(), f.name)
path_state = f.name
# Загрузить модель
loaded_model = Net()
loaded_model.load_state_dict(torch.load(path_state))
loaded_model.eval()
# Тестирование загруженной модели
x = torch.randn(2, 10)
output1 = model(x)
output2 = loaded_model(x)
print(Исходный вывод:, output1[0].tolist())
print(Вывод после загрузки:, output2[0].tolist())
# Очистка
os.remove(path_full)
os.remove(path_state)
Пример 4: Миграция устройства
Перенос модели между различными устройствами:
Пример
import torch
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc = nn.Linear(10, 5)
def forward(self, x):
return self.fc(x)
model = Net()
# Проверка текущего устройства
print(Устройство параметра:, model.fc.weight.device)
# Переместить на GPU (если доступно)
if torch.cuda.is_available():
model = model.cuda()
print(После перемещения на GPU:, model.fc.weight.device)
# Переместить обратно на CPU
model = model.cpu()
print(После перемещения обратно на CPU:, model.fc.weight.device)
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc = nn.Linear(10, 5)
def forward(self, x):
return self.fc(x)
model = Net()
# Проверка текущего устройства
print(Устройство параметра:, model.fc.weight.device)
# Переместить на GPU (если доступно)
if torch.cuda.is_available():
model = model.cuda()
print(После перемещения на GPU:, model.fc.weight.device)
# Переместить обратно на CPU
model = model.cpu()
print(После перемещения обратно на CPU:, model.fc.weight.device)
Пример 5: Инициализация с помощью apply
Используйте apply для рекурсивной инициализации:
Пример
import torch
import torch.nn as nn
def init_weights(module):
if isinstance(module, nn.Linear):
nn.init.xavier_uniform_(module.weight)
if module.bias is not None:
nn.init.zeros_(module.bias)
elif isinstance(module, nn.Conv2d):
nn.init.kaiming_normal_(module.weight)
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(10, 20)
self.fc2 = nn.Linear(20, 5)
def forward(self, x):
return self.fc2(self.fc1(x))
model = Net()
model.apply(init_weights)
print(fc1 вес:, model.fc1.weight[0, :3].tolist())
print(fc2 вес:, model.fc2.weight[0, :3].tolist())
import torch.nn as nn
def init_weights(module):
if isinstance(module, nn.Linear):
nn.init.xavier_uniform_(module.weight)
if module.bias is not None:
nn.init.zeros_(module.bias)
elif isinstance(module, nn.Conv2d):
nn.init.kaiming_normal_(module.weight)
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(10, 20)
self.fc2 = nn.Linear(20, 5)
def forward(self, x):
return self.fc2(self.fc1(x))
model = Net()
model.apply(init_weights)
print(fc1 вес:, model.fc1.weight[0, :3].tolist())
print(fc2 вес:, model.fc2.weight[0, :3].tolist())
Пример 6: Сложная структура сети
Постройте сеть с ветвями:
Пример
import torch
import torch.nn as nn
class BranchNet(nn.Module):
def __init__(self):
super(BranchNet, self).__init__()
# Основная часть
self.shared = nn.Linear(10, 20)
# Ветвь
self.branch_a = nn.Linear(20, 5)
self.branch_b = nn.Linear(20, 3)
def forward(self, x):
feat = self.shared(x)
out_a = self.branch_a(feat)
out_b = self.branch_b(feat)
return out_a, out_b
model = BranchNet()
x = torch.randn(4, 10)
out_a, out_b = model(x)
print(Выход ветви A:, out_a.shape)
print(Выход ветви B:, out_b.shape)
import torch.nn as nn
class BranchNet(nn.Module):
def __init__(self):
super(BranchNet, self).__init__()
# Основная часть
self.shared = nn.Linear(10, 20)
# Ветвь
self.branch_a = nn.Linear(20, 5)
self.branch_b = nn.Linear(20, 3)
def forward(self, x):
feat = self.shared(x)
out_a = self.branch_a(feat)
out_b = self.branch_b(feat)
return out_a, out_b
model = BranchNet()
x = torch.randn(4, 10)
out_a, out_b = model(x)
print(Выход ветви A:, out_a.shape)
print(Выход ветви B:, out_b.shape)
ЧАВО
Вопрос 1: Почему нужно вызывать super().__init__()?
Он вызывает метод инициализации родительского класса, обеспечивая корректную регистрацию параметров модулем.
<h3 Вопрос 2: Как просмотреть структуру модели?Непосредственно выведите print(model) или используйте функцию summary из torchvision.
<h3 Q3: В чем разница между parameters() и modules()?parameters() возвращает обучаемые параметры, тогда как modules() возвращает все модули.
Сценарии использования
nn.ModuleЭто основа для построения всех пользовательских нейронных сетей. Основные сценарии применения включают:
- Определение структуры сети: Любая пользовательская нейронная сеть
- Управление моделями: Сохранение, загрузка, миграция
- Управление параметрами: Обновления оптимизатора, вычисление градиентов
Примечание: Все пользовательские модули должны реализовывать метод forward.
Другие расширения