Метод parameters

Метод parameters класса Module возвращает итератор, который перебирает все обучаемые параметры (тензоры с requires_grad=True) модели и её подслоёв. Этот метод часто используется для передачи параметров модели в оптимизатор, а также для ручного обновления весов или применения к ним различных операций. Метод не принимает никаких параметров и возвращает объект-итератор.

Синтаксис

model.parameters()

Пример

Давайте создадим простую модель и получим её параметры:

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.fc = nn.Linear(5, 3)
		self.bn = nn.BatchNorm1d(3)
	
	def forward(self, x):
		return self.bn(self.fc(x))

model = SimpleModel()
params = list(model.parameters())
print(len(params))

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

4

Модель имеет четыре параметра: веса и смещение линейного слоя, а также параметры масштаба и смещения батч-нормализации.

Пример

Передадим параметры модели в оптимизатор SGD:

import torch
import torch.nn as nn
import torch.optim as optim

model = nn.Linear(10, 1)
optimizer = optim.SGD(model.parameters(), lr=0.01)
print(optimizer)

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

"SGD (Parameter Group 0: lr=0.01, momentum=0, dampening=0, weight_decay=0, nesterov=False)"

Пример

Параметры можно фильтровать, например, оставить только те, которые требуют градиент:

import torch
import torch.nn as nn

class CustomModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.linear = nn.Linear(5, 5)
		self.frozen = nn.Linear(5, 5)
		self.frozen.weight.requires_grad = False
	
	def forward(self, x):
		return self.linear(x) + self.frozen(x)

model = CustomModel()
trainable_params = [p for p in model.parameters() if p.requires_grad]
print(len(trainable_params))

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

2

В модели остались только два обучаемых параметра - веса и смещение слоя linear.

Пример

Применим операцию к каждому параметру модели, например, обнулим все градиенты вручную:

import torch
import torch.nn as nn

torch.manual_seed(0)
model = nn.Linear(5, 2)
loss = model(torch.randn(3, 5)).sum()
loss.backward()

print(model.weight.grad[0, 0])
for param in model.parameters():
	param.grad = None
print(model.weight.grad)

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

tensor(1.7000)
None

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

метод named_parameters,
который возвращает параметры вместе с их именами
метод state_dict,
который возвращает словарь всех параметров и буферов модели
метод load_state_dict,
который загружает параметры модели из словаря
метод zero_grad,
который обнуляет градиенты всех параметров модели