[PyTorch] nn.Module

https://pytorch.org/docs/stable/generated/torch.nn.Module.html

Module — PyTorch 1.13 documentation

---

named_children() & named_modules()

이 메서드들은 각각 모듈에 포함된 자식 모듈들과 모든 자식 모듈을 포함한 모듈의 계층 구조를 탐색할 수 있도록 합니다.

 

named_children()

각 자식 모듈의 이름과 해당 모듈 객체를 순회 가능한 iterator로 반환합니다.

즉, nn.Module 에서 정의한 모든 자식 모듈을 순회하면서 이름과 해당 모듈을 함께 사용할 수 있습니다.

import torch
import torch.nn as nn

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

model = MyModel()
for name, module in model.named_children():
    print(name, module)
    
# layer1 Linear(in_features=10, out_features=5, bias=True)
# layer2 Linear(in_features=5, out_features=2, bias=True)

named_modules()

모듈 계층 구조에서 자기 자신과 자신의 모든 하위 모듈을 모두 반복 가능한 iterator로 반홥합니다.

즉, 모델 내의 모든 모듈을 재귀적으로 탐색하면서 모듈 이름과 해당 모듈을 함께 사용할 수 있습니다.

import torch
import torch.nn as nn

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

model = MyModel()
for name, module in model.named_modules():
    print(name, module)
    
# MyModel(
#   (layer1): Linear(in_features=10, out_features=5, bias=True)
#   (layer2): Linear(in_features=5, out_features=2, bias=True)
# )
# layer1 Linear(in_features=10, out_features=5, bias=True)
# layer2 Linear(in_features=5, out_features=2, bias=True)

---

get_submodule()

주어진 문자열에 해당하는 하위 모듈을 반환하는 함수입니다. 주로 모델 내의 하위 모듈을 찾고자 할 때 사용됩니다.

다음과 같은 모델을 정의했을 때

import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(16)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(32)
        self.fc1 = nn.Linear(32*8*8, 10)

    def forward(self, x):
        x = nn.functional.relu(self.bn1(self.conv1(x)))
        x = nn.functional.relu(self.bn2(self.conv2(x)))
        x = x.view(-1, 32*8*8)
        x = self.fc1(x)
        return x

'conv1' 하위 모듈을 찾고싶으면

# MyModel 클래스의 인스턴스 생성
model = MyModel()

# `conv1` 하위 모듈 찾기
conv1 = model.get_submodule("conv1")

# 찾은 하위 모듈 출력하기
print(conv1)

# Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))

---

Parameter

named_parameters()

모델 내에 있는 파라미터들을 이름과 함께 딕셔너리 형태로 반환해주는 함수입니다.

import torch.nn as nn

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):
        x = self.fc1(x)
        x = self.fc2(x)
        return x

net = Net()
for name, param in net.named_parameters():
    print(name, param.shape)
    
# fc1.weight torch.Size([20, 10])
# fc1.bias torch.Size([20])
# fc2.weight torch.Size([5, 20])
# fc2.bias torch.Size([5])

 

get_parameter()

파라미터의 이름을 입력받아 해당 파라미터를 Tensor로 반환하는 함수입니다.

nn.Module 의 메서드로서 이 함수를 사용하려면 파라미터의 이름을 알아야 합니다.

import torch
import torch.nn as nn

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):
        x = self.fc1(x)
        x = self.fc2(x)
        return x

net = Net()
param = net.get_parameter("fc1.weight")
print(param)

# Parameter containing:
# tensor([[ 0.0547,  0.1041,  0.1740, -0.2452, -0.0330, -0.2526, -0.2220,  0.0566,
#          -0.0400,  0.2435],
#         ...
#         [ 0.1090,  0.2673, -0.0072,  0.0201, -0.2927, -0.2263, -0.1431, -0.1903,
#           0.3021,  0.1846]], requires_grad=True)

---

Buffer

named_buffers()

named_parameter() 메서드와 똑같은 기능을 합니다.

모델 내의 모든 buffer를 반복하면서 각 buffer의 이름과 값을 반환합니다.

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.bn = nn.BatchNorm1d(10)
    
    def forward(self, x):
        x = self.bn(x)
        return x

model = MyModel()

for name, buf in model.named_buffers():
    print(name, buf)
    
# bn.running_mean tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
# bn.running_var tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.])
# bn.num_batches_tracked tensor(0)

get_buffer()

모델 내의 특정 buffer를 가져올 때 사용됩니다.

buf = model.get_buffer('bn.running_mean')
print(buf)

# tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])

---

extra_repr()

모델의 print() 메서드에서 출력되는 문자열을 정의하는 방법 중 하나입니다. (str 형식으로 반환합니다.)

이 메서드는 모델의 파라미터 외에도 모델의 정보를 출력하는데 유용합니다.

import torch.nn as nn

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 2)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x
    
    def extra_repr(self):
        return "MyModule()"
    
model = MyModule()
print(model)

# MyModule(
#   MyModule()
#   (fc1): Linear(in_features=10, out_features=5, bias=True)
#   (fc2): Linear(in_features=5, out_features=2, bias=True)
# )

import torch.nn as nn

class Func_A(nn.Module):
    def __init__(self, name):
        super().__init__()
        self.name = name

    def forward(self, x):
        x = x * 2
        return x
    
    def extra_repr(self):
        return f'name={self.name}'

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.a = Func_A('function A 입니다.')
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 2)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x
    
    def extra_repr(self):
        return "MyModule()"
    
model = MyModule()
print(model)

# MyModule(
#   MyModule()
#   (a): Func_A(name=function A 입니다.)
#   (fc1): Linear(in_features=10, out_features=5, bias=True)
#   (fc2): Linear(in_features=5, out_features=2, bias=True)
# )

---

Docstring

Docstring을 추가할 수도 있습니다 !

from torch import nn

class Model(nn.Module):
    '''
    Docstring을 추가했습니다 !    
    '''
    def __init__(self):
        super().__init__()

model = Model()

print(model.__doc__)

#     Docstring을 추가했습니다 !

---

Hook

모델의 특정 레이어에서 입력, 출력, 혹은 그레이디언트 등을 중간에 가로 채어 원하는 작업을 수행하는 기능입니다.

모델의 레이어들은 forward()와 backward() 메소드를 호출하며, Hook을 사용하면 이러한 메소드가 호출되는 중간에 다양한 작업을 수행할 수 있습니다.

 

Foward hook

forward() 메소드 실행 시 호출되며, 입력, 출력값을 가로채서 수행합니다.

다음 모델의 두 번째 convolutional 레이어의 출력값 shape를 출력하는 Hook

import torch
import torch.nn as nn
import torch.nn.functional as F

# 모델 정의
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = x.view(-1, 9216)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 모델 생성
net = Net()

# Forward Hook 함수 정의
def print_output_shape(module, input, output):
    print("Output shape: ", output.shape)

# Forward Hook 등록
hook_handle = net.conv2.register_forward_hook(print_output_shape)

# 모델 실행
input = torch.randn(1, 1, 28, 28)
output = net(input)

# Output shape:  torch.Size([1, 64, 24, 24])

Backward Hook

backward() 메소드 실행 시 호출되며, 그레이디언트를 가로채서 수행합니다.

주로 gradient를 확인하거나 수정하는 데 사용됩니다.

import torch
import torch.nn as nn
import torch.nn.functional as F

# 모델 정의
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = x.view(-1, 9216)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 모델 생성
net = Net()

# backward hook 함수 정의
def backward_hook(module, grad_input, grad_output):
    print('grad_input:', grad_input)
    print('grad_output:', grad_output)

# backward hook 등록
net.conv1.register_backward_hook(backward_hook)

# 모델 실행
input = torch.randn(1, 1, 28, 28)
output = net(input)

# backward 연산 실행
output.sum().backward()

# grad_input: (None, tensor([[[[ 4.4821e-02,  7.7024e-03, -4.3794e-03],
#           ...
#           [-7.9876e-02,  2.8607e-04, -4.4666e-02]]]]), tensor([-0.0331, ... , 0.0140]))
# grad_output: (tensor([[[[ 0.0000e+00, -1.7648e-03,  0.0000e+00,  ...,  0.0000e+00,
#             0.0000e+00,  0.0000e+00],
#           ...
#           [ 0.0000e+00,  0.0000e+00, -1.0424e-03,  ...,  0.0000e+00,
#            -3.6263e-03,  0.0000e+00]]]]),)

---

Apply

모델의 모든 하위 모듈에 대해 지정된 함수를 적요하는 메서드입니다.

이 메서드는 모델을 변경하지 않고 하위 모듈의 파라미터나 버퍼 등에 접근하여 변경할 수 있는 유용한 방법입니다.

 

apply 메서드는 모델의 모든 하위 모듈에 대해 재귀적으로 실행됩니다. 각 하위 모듈에 대해 함수를 적용할 수 있으며, 이 함수는 하위 모듈을 입력으로 받아 처리하고 하위 모듈을 반환합니다.

 

예를 들어, 모든 선형 레이어의 가중치에 일괄적으로 초기화하려는 경우 다음과 같이 사용할 수 있습니다.

import torch.nn as nn

def init_weights(module):
    if isinstance(module, nn.Linear):
        module.weight.data.normal_(0.0, 0.02)

model = nn.Sequential(
    nn.Linear(1, 2),
    nn.ReLU(),
    nn.Linear(2, 3),
    nn.ReLU(),
    nn.Linear(3, 4),
    nn.ReLU()
)

print("가중치 변경 전 :")
for name, param in model.named_parameters():
    param = model.get_parameter(name)
    print(param)
print()

model.apply(init_weights)

print("가중치 변경 후 :")
for name, param in model.named_parameters():
    param = model.get_parameter(name)
    print(param)

또한, 다음과 같이 'requires_gard' 속성을 'False'로 변경할 수도 있습니다.

import torch.nn as nn

class MyModule(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(MyModule, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, output_size)
        
    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x
model = MyModule(1, 2, 5)

def set_parameters_no_grad(module):
    if hasattr(module, 'weight'):
        module.weight.requires_grad = False
    if hasattr(module, 'bias'):
        module.bias.requires_grad = False

print("속성 변경 전 :")
for name, params in model.named_parameters():
    param = model.get_parameter(name)
    print(param)
print()
        
model.apply(set_parameters_no_grad)

print("속성 변경 후 :")
for name, params in model.named_parameters():
    param = model.get_parameter(name)
    print(param)