Метод named_buffers

Метод named_buffers класса Module возвращает генератор, который перебирает все буферы модуля и всех его подслоев. Буферы - это тензоры, которые не являются параметрами модели и не обновляются во время обучения, но сохраняются в состоянии модели (например, средние значения и дисперсии в батч-нормализации). Метод возвращает кортежи вида (name, buffer), где name - это строка с полным именем буфера (с учётом иерархии подслоев), а buffer - сам тензор-буфер.

Синтаксис

module.named_buffers(prefix='', recurse=True, remove_duplicate=True)

Параметр prefix задаёт строку-префикс, которая будет добавлена к именам всех буферов. Параметр recurse управляет тем, будут ли включены буферы подслоев. Параметр remove_duplicate определяет, нужно ли удалять дублирующиеся буферы.

Пример

Создадим простую модель с буфером и получим все именованные буферы:

import torch
import torch.nn as nn

class MyModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.linear = nn.Linear(10, 5)
		self.register_buffer('running_mean', torch.zeros(5))
		
model = MyModel()
for name, buffer in model.named_buffers():
	print(f"{name}: {buffer.shape}")

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

"running_mean: torch.Size([5])"
"linear.weight: torch.Size([5, 10])"
"linear.bias: torch.Size([5])"

Обратите внимание, что метод named_buffers возвращает не только буферы, зарегистрированные через register_buffer, но и буферы подслоёв (например, у слоя Linear есть буферы весов и смещений).

Пример

Используем параметр prefix для добавления префикса к именам буферов:

import torch
import torch.nn as nn

class MyModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.register_buffer('running_mean', torch.zeros(3))
		
model = MyModel()
for name, buffer in model.named_buffers(prefix='model.'):
	print(name)

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

"model.running_mean"

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

Пример

Отключим рекурсивный обход подслоев с помощью параметра recurse:

import torch
import torch.nn as nn

class MyModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.linear = nn.Linear(2, 2)
		self.register_buffer('running_mean', torch.zeros(2))
		
model = MyModel()
for name, buffer in model.named_buffers(recurse=False):
	print(name)

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

"running_mean"

Как видим, буферы подслоя linear не были включены в вывод.

Пример

Удалим дублирующиеся буферы с помощью параметра remove_duplicate. Это может быть полезно, если один и тот же буфер используется в нескольких местах:

import torch
import torch.nn as nn

class MyModel(nn.Module):
	def __init__(self):
		super().__init__()
		buf = torch.ones(3)
		self.register_buffer('buf1', buf)
		self.register_buffer('buf2', buf)
		
model = MyModel()
print("C дубликатами:")
for name, _ in model.named_buffers(remove_duplicate=False):
	print(name)
print("Без дубликатов:")
for name, _ in model.named_buffers(remove_duplicate=True):
	print(name)

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

"C дубликатами:"
"buf1"
"buf2"
"Без дубликатов:"
"buf1"

По умолчанию дубликаты удаляются, поэтому повторяющиеся буферы будут возвращены только один раз.

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

метод buffers,
который возвращает только тензоры буферов без имён
метод parameters,
который возвращает параметры модели
метод named_parameters,
который возвращает именованные параметры модели
метод register_buffer,
который регистрирует новый буфер в модуле