[부스트캠프 AI Tech 5기] torch.nn

torch.nn
basic building block

torch.nn

nn.Linear

m = nn.Linear(20, 30)
input = torch.randn(128, 20)
output = m(input)
print(output.size())
torch.Size([128, 30])

• size 변환

nn.Identity

X = torch.Tensor([[1, 2],
                  [3, 4]])
m = nn.Identity()
m(X)
tensor([[1., 2.],
        [3., 4.]])

• 출력값이 동일

nn.LazyLinear

• Input size를 지정할 필요 없음

nn.Module

모델로 추상화 할 수 있는 클래스

• 빈 상자이므로 안에 nn.Module을 더 넣을 수 있음
• 계층적으로 담을 수 있음

nn.Sequential

model = nn.Sequential(
          nn.Conv2d(1,20,5),
          nn.ReLU(),
          nn.Conv2d(20,64,5),
          nn.ReLU()
        )

• 모듈을 하나로 묶어 순차적으로 실행

nn.ModuleList

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.linears = nn.ModuleList([nn.Linear(10, 10) for i in range(10)])

    def forward(self, x):
        # ModuleList can act as an iterable, or be indexed using ints
        for i, l in enumerate(self.linears):
            x = self.linears[i // 2](x) + l(x)
        return x

• 모듈을 리스트로 묶어 원하는 것을 실행

nn.ModuleDict

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.choices = nn.ModuleDict({
                'conv': nn.Conv2d(10, 10, 3),
                'pool': nn.MaxPool2d(3)
        })
        self.activations = nn.ModuleDict([
                ['lrelu', nn.LeakyReLU()],
                ['prelu', nn.PReLU()]
        ])

    def forward(self, x, choice, act):
        x = self.choices[choice](x)
        x = self.activations[act](x)
        return x

• 모듈을 딕셔너리로 묶어 원하는 것을 실행

정리

• Module은 서로 다른 Modele을 포함할 수 있음
• 최소 기능 단위인 function
• function으로 이루어진 layer
• layer로 이루어진 model
• 모든 딥러닝 모델은 이처럼 Module이 쌓여서 만들어짐

nn.Parameter

from torch.nn.parameter import Parameter
self.W = Parameter(torch.ones((out_features, in_features)))
self.b = Parameter(torch.ones(out_features))

• gradient 계산 가능
• 값 update
• 모델 저장시 값 저장

Buffer

self.register_buffer('name', tensor)

• 모델 저장시 값 저장

nn.Module 분석

model.named_modules()

• 모델에 속하는 모든 submodule 표시

model.named_children()

• 한 단계 아래의 submodule 만 표시

model.get_submodule()

model.get_submodule('path')

• 모델의 layer name을 따라 path 로 지정 ab.a.c 요로코롬

model.get_parameter()

model.get_parameter('path')

• 모델의 layer name을 따라 path 로 지정 ab.b.W1 요로코롬

model.get_buffer()

model.get_buffer('path')

• 모델의 layer name을 따라 path 로 지정 cd.c.duck 요로코롬

• 가장 낮은 단위인 module간 참조는 불필요

extra_repr

class Function_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 str('name='+self.name)

• extra_repr를 재정의 해줘야 원하는 출력이 나옴

hook

package 중간중간에 custom code를 실행할 수 있도록 미리 만들어놓은 인터페이스

model.register_forward_pre_hook(hook)
model.register_forward_hook(hook)
model.register_full_backward_hook(hook)
model.W.register_hook(tensor_hook)

• forward pre hook
• forward hook
• backward hook
• backward tensor hook

apply

module을 입력으로 받아 모든 module에 함수를 적용시켜주는 코드
일반적으로 가중치 초기화에 많이 사용

# apply가 적용된 module을 return 
returned_module = model.apply(print_module)