[pytorch] nn.BatchNorm 알아보기

 

- 목차

 

키워드.

• - Batch Normalization

 

들어가며.

이번 글에서는 pytorch 의 BatchNorm 에 대해서 알아보려고 합니다.

Batch Normalization 의 개념과 여러가지 사례들, 그리고 pytorch.nn.BatchNorm 의 활용법에 대해서 알아봅니다.

 

Batch Normlization 이란 ?

Batch Normalization 은 이름 그대로 Batch Data 를 대상으로 Normalization 을 적용합니다.

여기서 Batch Data 는 Neural Network 로 순전파되는 Batch Size 의 입력 데이터를 의미합니다.

pytorch 를 예로 들면, DataLoader 의 batch_size 에 설정된 Batch Size 를 의미합니다.

아래의 이미지 예시와 같이 총 764 개의 이미지 중에서 Batch Size 를 4개로 설정하게 되면, 

4개의 데이터를 한 묶음으로 하는 Batch Data, Batch Tensor 들로 나뉘게 됩니다.

출처 : https://www.kaggle.com/datasets/alxmamaev/flowers-recognition

 

 

그리고 Normalization 은 Standard Normalization 을 의미합니다.

0 을 중심으로 삼고 Batch Data 의 평균과 분산을 기준으로 하는 Normal Distribution 을 만들게 됩니다.

아래와 같이 feature 가 1개인 데이터에 대해서 6개의 batch 를 6x1 사이즈의 Tensor 로 만들었다고 가정하겠습니다.

이를 Standard Normalization 하게 되면, Mean 과 Variance 를 고려하여 다음과 같이 변환되게 됩니다.

참고로 Standard Normalization 의 수식은 $$ \frac{x - \mu}{\delta} $$ 와 같습니다. ( $\mu$ 는 평균, $\delta$ 는 표준편차 )

그리고 Standard Normalization 을 적용한 결과는 아래와 같습니다.

 

이러한 방식으로 Neural Net 의 Forward Propagation 단계에서 Batch Data 를 대상으로 Standard Normalization 이 이루어집니다.

그리고 정규화된 데이터는 다음 레이어로 전달됩니다.

 

nn.BatchNorm 사용해보기.

이제 pytorch 의 nn.BatchNorm 레이어를 만들어 Batch Normalization 을 적용해보도록 하겠습니다.

Batch Normalization 은 입력 데이터의 차원에 따라서 1d, 2d 등의 모듈이 제공됩니다.

이번 글에서는 nn.BatchNorm1d 에 대해서 알아봅니다.

nn.BatchNorm1d  Arguments 알아보기.

먼저 nn.BatchNorm1d 의 인자에 대해서 알아보겠습니다.

 

num_features.

num_features 는 입력 데이터의 feature 의 갯수를 의미합니다.

BatchNorm1d 는 입력 데이터의 feature 갯수만큼 Batch Normalization 을 수행합니다.

그래서 nn.Linear 가 in_features, out_features 를 가지듯 feature 의 갯수를 입력받아야하며, 그 인자의 이름이 num_features 입니다.

 

아래의 코드예시는 num_features 이 1 일때의 BatchNorm1d 의 파라미터를 보여줍니다.

결과적으로 1개의 Weight 와 1개의 Bias 를 가집니다.

Batch Normalizaion 은 그 내부적으로 $\gamma$ 와 $\beta$ 라는 2개의 파라미터를 가지는데요.

$\gamma$ 와 $\beta$ 는 Batch Normalization 을 적용한 Neural Network 에서 학습해야한 파라미터입니다.

이에 대한 자세한 얘기는 Batch Normalization 의 개론적인 설명을 위한 새로운 글로 설명드리도록 하겠습니다.

import torch
import torch.nn as nn
bn = nn.BatchNorm1d(num_features=1)
for parameter in bn.parameters():
    print(parameter)

Parameter containing:
tensor([1.], requires_grad=True)
Parameter containing:
tensor([0.], requires_grad=True)

 

아래의 예시는 입력 데이터가 10개의 Feature 를 가지는 경우에, Batch Norm 레이어의 Parameter 를 보여줍니다.

10개의 Feature 에 대해서 각각 Normalization 을 적용해야하기 때문에 총 10개의 Weight 와 10개의 Bias 를 가집니다.

import torch
import torch.nn as nn
bn = nn.BatchNorm1d(num_features=10)
for parameter in bn.parameters():
    print(parameter)

Parameter containing:
tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.], requires_grad=True)
Parameter containing:
tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], requires_grad=True)

 

momentum.

Batch Normalization 은 Batch Data 를 대상으로 정규화시킨다고 말씀드렸죠 ?

일반적인 Deep Learning 은 Total Data 를 여러개의 Batch Data 로 나누어서 학습을 진행합니다.

즉, Total Data 가 100개이고, Batch Data 의 크기가 10 개라면 한번의 학습에서 총 10번의 Iteration 이 발생합니다.

그럼 순전파 과정에서 10번의 Batch Normalization 이 발생하게 되죠. ( 왜냐하면 Batch Data 가 10개이기 때문에 )

 

이 과정에서 10번의 Forward/Back Propagation 이 발생하게 되고, BatchNorm1d 의 Weight 와 Bias 는 10번 Update 가 됩니다.

BatchNorm1d 의 Weight 와 Bias 는 학습이 거듭될수록 그 값이 더 정교해지겠죠 ?

Initial Weight 와 Bias 는 1 과 0 에서 최적의 값으로 수정되어 갈테니까요.

 

그리고 모든 학습이 마무리가 되어, Testing 을 수행하거나 실제 서비스 단계에서 활용되는 상황이 됩니다.

이때에는 Training 단계에서 Mini-Batch 를 대상으로 계산했던 각 Batch 의 평균과 분산을 토대로 마지막 평균과 분산을 생성합니다.

(만약 10번의 Batch Data 들을 학습했다면, 10개의 Batch Mean 과 10개의 Batch Variance 가 존재할 겁니다. )

그리고 테스트나 서비스 운영 단계에서는 학습 단계에서 생성된 모든 Batch Mean 과 Batch Variance 를 평균을 내어  최종적인 Mean 과 Variacen 를 계산합니다.

이 과정에서 학습 초기의 Mean 과 Variance 에 적은 중요도를 부여하고, 학습 마무리 단계에서 얻은 결과값을 Mean 과 Variance 에 높은 가중치를 부여하는 방식이 Momentum 입니다.

 

 

nn.BatchNorm1d 간단한 예시.

feature 를 1개를 가지는 Tensor 에 대해서 Batch Normalization 을 적용해보겠습니다.

아래의 예시에서는 1 ~ 6 까지의 자연수들을 하나의 Batch 로 묶어서 Batch Norm Layer 를 통과시킵니다.

batch size 는 6이고, 평균은 2.5, 분산과 표준편차는 각각 2.91, 1.7 를 가집니다.

import torch
import torch.nn as nn
bn = nn.BatchNorm1d(num_features=1)
tensor = torch.tensor([1,2,3,4,5,6], dtype=torch.float32).view(-1, 1)
bn(tensor)

tensor([[-1.4638],
        [-0.8783],
        [-0.2928],
        [ 0.2928],
        [ 0.8783],
        [ 1.4638]], grad_fn=<NativeBatchNormBackward0>)

 

 

 

Batch Norm 를 통해서 Vanishing Gradient 문제 해결하기.

먼저 Batch Normalization 을 활용하여 Dying ReLU 현상을 극복하는 예시 코드를 살펴보겠습니다.

아래의 모델은 Linear Layer 를 활용한 Linear Regressor 모델이구요.

Dying ReLU 현상을 재현하기 위해서 극단적인 모델을 만들었습니다.

모델은 단순히 Linear Layer 와 ReLU 로 구성되어 있으며, Input Tensor 에 해당하는 데이터는

$$ y = 1 \times x $$ 를 구성하는 데이터들로 준비하였습니다.

 

하지만 Linear Layer 의 Weight 를 -10 으로 설정하였습니다.

Linear Layer 의 Weight 가 -10 인 경우에는 ReLU Activation Function 에 의해서 Gradient 가 0 이므로 이는 정상적인 Back Propagation 이 불가능합니다.

import torch
import torch.nn as nn

class TestModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1, bias=False)
        self.linear.weight = nn.Parameter(torch.tensor([[-10.0]], requires_grad=True))
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(self.linear(x))
        
model = TestModel()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
loss_function = nn.MSELoss()

X_train = torch.linspace(1, 1000, 10000, dtype=torch.float32).reshape((-1,1))
y_train = torch.linspace(1, 1000, 10000, dtype=torch.float32).reshape((-1,1))

all_losses = []
for epoch in range(10000):
    output = model(X_train)
    loss = loss_function(output, y_train)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    all_losses.append(loss.item())
    print(epoch, loss.item())

import matplotlib.pyplot as plt
plt.plot(all_losses)
plt.show()

 

 

위 데이터셋과 모델의 학습 결과로써 얻을 수 있는 Loss Graph 는 아래와 같이, 0 으로 접근하지 못합니다.

 

여기서 Batch Normalization 을 적용해보겠습니다.

 

import torch
import torch.nn as nn

class TestModelBN(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1, bias=False)
        self.linear.weight = nn.Parameter(torch.tensor([[-10.0]], requires_grad=True))
        self.relu = nn.ReLU()
        self.bn = nn.BatchNorm1d(1)

    def forward(self, x):
        return self.relu(self.bn(self.linear(x)))

model = TestModelBN()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
loss_function = nn.MSELoss()

X_train = torch.linspace(1, 1000, 10000, dtype=torch.float32).reshape((-1,1))
y_train = torch.linspace(1, 1000, 10000, dtype=torch.float32).reshape((-1,1))

all_losses = []
for epoch in range(10000):
    output = model(X_train)
    loss = loss_function(output, y_train)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    all_losses.append(loss.item())
    print(epoch, loss.item())

import matplotlib.pyplot as plt
plt.plot(all_losses)
plt.show()

 

Batch Normalization 을 적용하게 되면 위와 같이 Loss Graph 가 0 으로 수렴하게 됩니다.

그 이유는 다음과 같습니다.

아래의 데이터는 각 Layer 를 통과한 이후의 데이터 모습을 보여줍니다.

1, 2, 3, 4, 5 에 해당하는 Train Data 는 Linear Layer 를 지나면서 -10 배만큼 값이 증가하게 되는데요.

이는 ReLU 함수를 지나게 되면서 0 으로 값이 변경되게 됩니다.

그리고 필연적으로 역전파로 인한 파라미터 업데이트가 불가능해지므로 비효율적인 모델이 됩니다.

epoch 1
training data    :  [1.0, 2.0, 3.0, 4.0, 5.0]
after Linear data:  [-10.0, -20.0, -30.0, -40.0, -50.0]
after ReLU   data:  [0.0, 0.0, 0.0, 0.0, 0.0]

epoch 10
training data    :  [1.0, 2.0, 3.0, 4.0, 5.0]
after Linear data:  [-10.0, -20.0, -30.0, -40.0, -50.0]
after ReLU   data:  [0.0, 0.0, 0.0, 0.0, 0.0]

epoch 100
training data    :  [1.0, 2.0, 3.0, 4.0, 5.0]
after Linear data:  [-10.0, -20.0, -30.0, -40.0, -50.0]
after ReLU   data:  [0.0, 0.0, 0.0, 0.0, 0.0]

epoch 1000
training data    :  [1.0, 2.0, 3.0, 4.0, 5.0]
after Linear data:  [-10.0, -20.0, -30.0, -40.0, -50.0]
after ReLU   data:  [0.0, 0.0, 0.0, 0.0, 0.0]

epoch 5000
training data    :  [1.0, 2.0, 3.0, 4.0, 5.0]
after Linear data:  [-10.0, -20.0, -30.0, -40.0, -50.0]
after ReLU   data:  [0.0, 0.0, 0.0, 0.0, 0.0]

epoch 9999
training data    :  [1.0, 2.0, 3.0, 4.0, 5.0]
after Linear data:  [-10.0, -20.0, -30.0, -40.0, -50.0]
after ReLU   data:  [0.0, 0.0, 0.0, 0.0, 0.0]

 

반면 아래의 로그는 Batch Normalization 을 적용한 모델의 로그입니다.

아래와 같이 epoch 을 거듭할수록 Batch Normalization 에 의해서 -10 배만큼 증가된 값들이 정규화되고,

정규화된 결과로 인해서 정상적인 학습이 가능해집니다.

epoch 1
training data    :  [1.0, 2.0, 3.0, 4.0, 5.0]
after Linear data:  [-10.0, -20.0, -30.0, -40.0, -50.0]
after BN     data:  [1.4148, 0.7076, 0.0004, -0.7068, -1.414]
after ReLU   data:  [1.4148, 0.7076, 0.0004, 0.0, 0.0]

epoch 10
training data    :  [1.0, 2.0, 3.0, 4.0, 5.0]
after Linear data:  [-10.0, -20.0, -30.0, -40.0, -50.0]
after BN     data:  [1.4302, 0.7222, 0.0142, -0.6938, -1.4018]
after ReLU   data:  [1.4302, 0.7222, 0.0142, 0.0, 0.0]

epoch 100
training data    :  [1.0, 2.0, 3.0, 4.0, 5.0]
after Linear data:  [-10.0, -20.0, -30.0, -40.0, -50.0]
after BN     data:  [1.5683, 0.8566, 0.1449, -0.5669, -1.2786]
after ReLU   data:  [1.5683, 0.8566, 0.1449, 0.0, 0.0]

epoch 1000
training data    :  [1.0, 2.0, 3.0, 4.0, 5.0]
after Linear data:  [-10.0, -20.0, -30.0, -40.0, -50.0]
after BN     data:  [1.5555, 1.703, 1.8506, 1.9982, 2.1458]
after ReLU   data:  [1.5555, 1.703, 1.8506, 1.9982, 2.1458]

epoch 5000
training data    :  [1.0, 2.0, 3.0, 4.0, 5.0]
after Linear data:  [-10.0, -20.0, -30.0, -40.0, -50.0]
after BN     data:  [1.0002, 1.9999, 2.9996, 3.9993, 4.999]
after ReLU   data:  [1.0002, 1.9999, 2.9996, 3.9993, 4.999]

epoch 9999
training data    :  [1.0, 2.0, 3.0, 4.0, 5.0]
after Linear data:  [-10.0, -20.0, -30.0, -40.0, -50.0]
after BN     data:  [1.0, 2.0, 2.9999, 3.9999, 4.9999]
after ReLU   data:  [1.0, 2.0, 2.9999, 3.9999, 4.9999]

 

 

사실상 위에서 설명한 내용은 매우 극단적인 케이스에 해당합니다.

Weight Initialization 만 잘해도 Dying ReLU 나 Vanishing Gradient 문제에 직면하지 않을 수 있습니다.

하지만, Signoid 나 Tanh 함수는 깊은 모델에서 사용한다면 Vanishing Gradient 문제를 해결하지 못할 수 있습니다.

이러한 경우에서 Batch Normalization 이 큰 위력을 발휘할 수 있습니다.