Updated:

PyTorch Documentation: Base class for all neural network modules.

import torch
import torch.nn as nn

01. torch.nn

신경망(Neural Network)를 만들기 위한 layer, cost function을 모아놓은 package. PyTorch Docs에는 torch.nn을 the basic buliding block이라고 표현했을만큼 사실상 신경망 모델을 구축하는데 있어 필요한 모든 것을 갖추고 있다고 볼 수 있다.

torch.nn은 Layer, activation function, loss function 등 모델을 구축하는 데 있어 필요한 layer, function과 이러한 layer, function을 하나로 모아서 신경망을 만들 수 있는 Container로 구성이 되어있다.

이번 글에서는 torch.nn 중에서도 Module이라는 class를 다룰 예정이다. 전체적인 구조와 attribute, method들이 어떤 것이 존재하는 지 살펴보고 간단한 예제와 함께 그 기능을 익힐 것이다.

참고: 일반적으로 torch.nn 패키지는 nn 이라는 alias로 많이 import한다.

02. nn.Module

PyTorch로 신경망을 구현한 모든 class의 부모가 되는 class. PyTorch Docs에서는 Base class라고 표현되어 있을만큼 중요한 class이다.

앞으로 직접 만들어보게 될 신경망도 PyTorch로 만들게된다면 역시 Module class를 상속받아서 구현해야한다.

Module은 layer, function, 다른 신경망 모델 모두 포함할 수 있고 트리 구조 형태로 포함하고 있는 모델, layer, function을 구축한다. 보통은 attribute에 layer, function, 모델을 받아서 만든다.

2.1. nn.Module를 상속받아 기초적인 모델 구현해보기

아래는 nn.Module을 상속받아 $y=ax+b$를 구현한 간단한 Linear Regression 모델이다.

class LinearModel(nn.Module):
    def __init__(self):
        super(LinearModel, self).__init__()
        self.weight = nn.Parameter(torch.rand(1))
        self.bias = nn.Parameter(torch.rand(1))

    def forward(self,x):
        return x * self.weight + self.bias

nn.Module을 상속받아 신경망을 만들 때는 가장 기본적으로 init과 forward가 구현되어 있어야한다.

init

init에는 신경망에 필요한 parameter를 nn.Parameter를 이용하여 정의한다. nn.Parameter는 신경망에 필요한 parameter를 만들어주는 class이다. 별다른 keyword를 추가하지 않아도 자동적으로 미분 계산 값을 위한 keyword 인 requires_grad=True를 포함시켜준다. 그리고 nn.Parameter로 지정된 parameter는 신경망 저장시 같이 저장된다. 신경망의 parameter를 일반 Tensor로 만들게 정의하게 되면 미분 계산을 할 수 있을 지 몰라도 모델 저장시 같이 저장이 되지 않는다.

forward

forward는 예측값을 반환하는 메소드이다. 주어진 예시는 Linear Regression이기 때문에 위와 같이 예측값을 만들어 반환하도록 구현하였다.

2.2. parameter 관련 메소드

nn.Parameter와 관련된 메소드가 아닌 신경망이 지닌 parameter를 반환하는 메소드이다.

parameters()

  • Module이 가진 parameter에 대한 iterator를 반환한다.

named_parameters()

  • Module이 가진 parameter에 대한 iterator를 반환한다.
  • Module.parameters()와의 차이점: name과 parameter를 같이 반환한다.

get_parameter(target)

  • target으로 지정된 parameter 반환
  • 없으면 AttributeError 발생
linear_model = LinearModel()

print("-----------parameters()------------")
for param in linear_model.parameters():
    print(param)

print("-----------named_parameters()-------")
for name, param in linear_model.named_parameters():
    print(f"name: {name} \nparams: {param}")

print("-----------get_parameter()------------")
print(linear_model.get_parameter('weight'))
-----------parameters()------------
Parameter containing:
tensor([0.7849], requires_grad=True)
Parameter containing:
tensor([0.4208], requires_grad=True)
-----------named_parameters()-------
name: weight
params: Parameter containing:
tensor([0.7849], requires_grad=True)
name: bias
params: Parameter containing:
tensor([0.4208], requires_grad=True)
-----------get_parameter()------------
Parameter containing:
tensor([0.7849], requires_grad=True)

2.3. buffer관련 메소드

buffer는 parameter는 아니지만 신경망에 등록된 Tensor를 가리킨다. 여기서 등록되었다는 의미는 신경망을 저장할 때 같이 저장된다는 뜻이다. buffer는 parameter가 아니기 때문에 gradient를 필요로 하지 않는다. 보통 특정 layer의 mean과 std를 저장하는데 쓰인다.

register_buffer(name, tensor)

  • tensor를 name이라는 이름으로 신경망에 등록시키는 함수

참고: what pytorch means by buffers

linear_model = LinearModel()
linear_model.register_buffer("buf",torch.tensor([4]))

linear_model.buf
tensor([4])

buffers()

  • 신경망이 가진 buffer를 iterator로 반환해주는 함수

named_buffers()

  • 신경망이 가진 buffer를 iterator로 반환해주는 함수
  • buffers()와의 차이점: 이름을 반환해준다.

get_buffer(target)

  • target으로 등록된 buffer 반환.
  • 없으면 AttributeError 발생
print("-----------buffer()------------")
for buffer in linear_model.buffers():
    print(buffer)

print("-----------named_buffers()------")
for name, buffer in linear_model.named_buffers():
    print(f"name: {name}\nbuffer: {buffer}")

print("-----------get_buffer()--------")
print(linear_model.get_buffer('buf'))
-----------buffer()------------
tensor([4])
-----------named_buffers()------
name: buf
buffer: tensor([4])
-----------get_buffer()--------
tensor([4])

2.4 module, submodule 관련 함수

modules()

  • 신경망이 가진 module들을 iterator로 반환

named_modules()

  • 신경망이 가진 module들을 이름과 함께 iterator로 반환

children()

  • 해당 신경망의 자식 신경망만 반환

named_children()

  • 해당 신경망의 자식 신경망과 이름을 반환

get_submodule(target)

  • target으로 지정된 submodule을 반환한다.
  • 없으면 Attribute Error 발생
linear_model1 = LinearModel()
linear_model2 = LinearModel()
linear_model3 = LinearModel()

linear_model1.sub1 = linear_model2
linear_model2.sub2_1 = linear_model3

print("-----------module()------------")
for module in linear_model1.modules():
    print(module)

print("-----------named_module()------------")
for name,module in linear_model1.named_modules():
    print(f"name: {name}\nmodule: {module}")

print("-----------children()------------")
for module in linear_model1.children():
    print(module)
print("-----------named_children()------------")
for name, module in linear_model1.named_children():
    print(f"name: {name}\nmodule: {module}")

print("-----------get_submodule()------------")
print(linear_model1.get_submodule('sub1'))
-----------module()------------
LinearModel(
  (sub1): LinearModel(
    (sub2_1): LinearModel()
  )
)
LinearModel(
  (sub2_1): LinearModel()
)
LinearModel()
-----------named_module()------------
name:
module: LinearModel(
  (sub1): LinearModel(
    (sub2_1): LinearModel()
  )
)
name: sub1
module: LinearModel(
  (sub2_1): LinearModel()
)
name: sub1.sub2_1
module: LinearModel()
-----------children()------------
LinearModel(
  (sub2_1): LinearModel()
)
-----------named_children()------------
name: sub1
module: LinearModel(
  (sub2_1): LinearModel()
)
-----------get_submodule()------------
LinearModel(
  (sub2_1): LinearModel()
)

2.5. apply

apply(fn)

  • fn: 함수
  • 신경망에 등록된 모든 submodule에 fn을 적용한다.
def make_new_weight(module):
    module.weight2 = nn.Parameter(torch.tensor([2.0]))


linear_model1 = LinearModel()
linear_model2 = LinearModel()

linear_model1.sub = linear_model2

linear_model1.apply(make_new_weight)

print(linear_model1.get_parameter("weight2"))
print(linear_model2.get_parameter("weight2"))
Parameter containing:
tensor([2.], requires_grad=True)
Parameter containing:
tensor([2.], requires_grad=True)

3. nn.Module을 이용한 신경망 학습 과정

위에서 간단히 구현한 LinearModel class를 이용하여 nn.Module을 구현한 신경망의 학습이 어떻게 진행되는지 간단하게 살펴본다.

$y=2x+1$ 을 학습한다고 가정해보자.

class LinearModel(nn.Module):
    def __init__(self):
        super(LinearModel, self).__init__()
        self.weight = nn.Parameter(torch.rand(1))
        self.bias = nn.Parameter(torch.rand(1))

    def forward(self,x):
        return x * self.weight + self.bias

먼저 x에 학습 데이터, y에는 실제값을 지정한다.

# 학습 데이터 지정, 실제값 지정
# 학습 데이터의 (data 수, feature 수) shape을 가진 2D-Tensor로 만들어야한다.
x = torch.arange(10).float().view(-1,1)
y = 2*x + 1

모델 생성을 생성하고 loss funciton, loss function을 최소화할 parameter를 찾는 optimizer를 지정한다.

  • opimizer: SGD(stochastic gradient descent, 확률적 경사하강법)
    • optimizer에 갱신할 parameter를 지정해주어야한다.
  • loss function: MSE(mean squared error), 회귀 문제에서 일반적으로 쓰인다.
import torch.optim as optim

linear_model = LinearModel()

# optimizer에 갱신할 parameter 지정, 학습률은 0.01 값 사용
optimizer = optim.SGD(linear_model.parameters(), lr=0.01)
criterion = nn.MSELoss()

학습은 다음과 같이 이루어진다.

학습과정

  1. model을 통해 주어진 data의 예측값 반환
  2. 예측값과 실제값을 이용해 loss를 구한다.
  3. loss를 이용해 parameter의 gradient값을 구한다.
  4. 구한 gradient 값을 이용해 parameter를 갱신한다.
  5. 지정된 학습 횟수만큼 반복

예측값 반환 방법

예측값을 반환하는 방법은 2가지가 있다.

  • 객체를 호출해 예측값을 구하는 방법 e.g) linear_model(x)
  • 모델이 가진 forward() 메소드 사용 e.g.) linear_model.forward(x)

보통 위의 방법을 사용한다.

# 학습 횟수
EPOCH = 10

for epoch in range(EPOCH):
    # 예측값 반환
    pred = linear_model(x)

    # loss 구함
    loss = criterion(pred,y)

    # 각 parameter의 grad 초기화
    optimizer.zero_grad()

    # 각 parameter의 gradient 구함
    loss.backward()

    # parameter 갱신
    optimizer.step()
    print(f"EPOCH {epoch} COST: {loss.item()}")
EPOCH 0 COST: 37.13650131225586
EPOCH 1 COST: 6.440334320068359
EPOCH 2 COST: 1.1369059085845947
EPOCH 3 COST: 0.22039809823036194
EPOCH 4 COST: 0.061789851635694504
EPOCH 5 COST: 0.0341210812330246
EPOCH 6 COST: 0.029076844453811646
EPOCH 7 COST: 0.027944350615143776
EPOCH 8 COST: 0.02749050036072731
EPOCH 9 COST: 0.02715681865811348
# 학습 결과를 바탕으로 예측
pred = linear_model(x)

# 1D-Tensor로 출력
print(f"x:{x.flatten()}")
print(f"예측값: {pred.flatten()}")
print(f"y: {y.flatten()}")
x:tensor([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])
예측값: tensor([ 0.6955,  2.7439,  4.7923,  6.8406,  8.8890, 10.9374, 12.9858, 15.0342,
        17.0826, 19.1310], grad_fn=<ReshapeAliasBackward0>)
y: tensor([ 1.,  3.,  5.,  7.,  9., 11., 13., 15., 17., 19.])
test = torch.tensor([[10.0]])
pred_t = linear_model(test)
print(f"예측값: {pred_t}")
예측값: tensor([[21.1794]], grad_fn=<AddBackward0>)

Leave a comment