Метод get_submodule

Метод get_submodule класса Module возвращает вложенный модуль по указанному строковому пути. Первым параметром метод принимает строку с путём к модулю, где части пути разделяются точками. Если модуль по указанному пути не найден, метод выбрасывает исключение AttributeError. Метод удобен для получения доступа к глубоко вложенным подслоям модели без необходимости разбирать структуру вручную.

Синтаксис

module.get_submodule(target)

Пример

Давайте создадим простую модель и получим доступ к вложенному линейному слою по его имени:

import torch
import torch.nn as nn

class MyModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.layer1 = nn.Linear(10, 20)
		self.layer2 = nn.Linear(20, 5)

model = MyModel()
submodule = model.get_submodule('layer1')
print(submodule)

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

Linear(in_features=10, out_features=20, bias=True)

Пример

Давайте создадим модель с несколькими уровнями вложенности и получим доступ к модулю по составному пути:

import torch
import torch.nn as nn

class NestedBlock(nn.Module):
	def __init__(self):
		super().__init__()
		self.fc = nn.Linear(5, 3)

class DeepModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.block1 = NestedBlock()
		self.block2 = NestedBlock()

model = DeepModel()
submodule = model.get_submodule('block1.fc')
print(submodule)

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

Linear(in_features=5, out_features=3, bias=True)

Пример

Давайте попробуем получить доступ к несуществующему модулю. Метод выбросит исключение:

import torch
import torch.nn as nn

class MyModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.layer1 = nn.Linear(10, 20)

model = MyModel()
try:
	submodule = model.get_submodule('layer2')
except AttributeError as e:
	print(e)

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

"has no attribute 'layer2'"

Пример

Давайте применим метод к модели Sequential и получим доступ к одному из слоёв по индексу:

import torch
import torch.nn as nn

model = nn.Sequential(
	nn.Linear(10, 20),
	nn.ReLU(),
	nn.Linear(20, 5)
)
submodule = model.get_submodule('2')
print(submodule)

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

Linear(in_features=20, out_features=5, bias=True)

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

метод add_module,
который добавляет дочерний модуль в текущий
метод get_parameter,
который возвращает параметр по имени
метод get_buffer,
который возвращает буфер по имени
метод named_modules,
который возвращает все дочерние модули с их именами