본문 바로가기

책/파이토치 트랜스포머를 활용한 자연어 처리와 컴퓨터비전 심층학습

11 모델 배포 (2) 모델 경량화

1. Knowledge Distillation

지식 증류(Knowledge Distillation)는 복잡한 신경망은 교사 신경망의 지식을 학습하는 단순한 신경망인 학생 신경망에 전달하여 성능을 개선하는 방법이다.

 

 

응답 기반 지식 증류(Response-based)

교사 신경망의 출력값을 활용해 학생 신경망을 학습시키는 방법이다. 교사 신경망의 출력값을 soft label(각 클래스에 대한 확률분포)로 사용하요 학생 신경망이 해당 확률 분포를 모방하도록 학습한다.

 

확률분포는 소프트맥스로 구하는데 이때 온도 매개변수 T를 추가한다. T는 소프트맥스 출력 분포를 변화시키는 용도인데, T가 크면 출력 분포가 넓어지고, T가 작으면 출력 분포가 좁아진다.

 

 

손실 함수를 코드로 보면 다음과 같다.

def distillation_loss(student_output, labels, teacher_output, T, alpha):
  student_softmax = F.log_softmax(student_output / T, dim=1)
  teacher_softmax = F.softmax(teacher_output / T, dim=1)
  temperature_loss = T * T * 2.0 * alpha

  kld_loss = nn.KLDivLoss(reduction='batchmean')(student_softmax, teacher_softmax)
  kld_loss *= temperature_loss
  ce_loss = F.cross_entropy(student_output, labels) * (1.0 - alpha)

  return kld_loss + ce_loss

 

교사와 학생의 소프트맥스를 구할 때 출력값을 T로 나눈뒤 소프트맥스 함수를 적용한다. 이때 학생의 경우 log_softmax를 사용하는 데, 이는 계산을 안정화하기 위해서이다. 만약 교사도 log를 취한다면 nn.KLDivLoss의 log_target 매개변수(default=False)를 True로 지정해주어야 한다.

 

쿨백 라이블러 발산 손실과 학생과 label의 교차 엔트로피 손실을 더해서 최종 손실로 사용한다. 최종 손실을 계산할 때는 alpha 매개변수로 두 손실의 비율을 조정하여 합한다.

 

 

특징 기반 지식 증류(Feature-based)

교사 신경망의 중간층의 출력값을 학생 신경망의 중간층에 전달하여 학생 신경망이 교사 신경망의 특징을 학습하도록 유도한다.

 

 

관계 기반 지식 증류(Relation-based)

신경망에서 학습된 다른 층 간의 관계를 이용하여 지식을 증류하는 방법이다. 손실 함수는 다음과 같이 표현할 수 있다.

 

$L_{RelD}(f_t, f_s) = L_R(\psi(f_t), \psi(f_s))$

 

관계 함수 $\psi$는 신경망의 서로 다른 층간 관계 정보를 계산하는 함수로 Gram Matrix를 이용한 방법이 주로 사용된다. $L_R$은 교사 신경망과 학생 신경망의 관계정보의 유사도를 계산하는 함수로, 두 신경망의 관계 정보가 최대한 유사하게 만든다.

 

어떤 벡터들의 집합 V의 Gram matrix는 다음과 같이 계산된다.

 

$G = V^T V$

 

그림으로 표현하면 다음과 같다.

 

 

2. 텐서 분해

텐서 분해는 다차원 텐서를 여러 개의 작은 텐서로 분해하는 기법이다. 인공 신경망의 가중치 텐서를 분해한다면 더 효율적으로 저장하고 계산 효율성을 향상시킬 수 있다.

 

인공 신경망에서는 텐서 분해의 한 형태인 Low-Rank Decomposition를 적용한다. 이는 다차원 텐서 $M$를 저차원 텐서 $\hat M$로 근사하는 기법이다.

 

주로 2차원 행렬은 SVD로 분해하고, 3차원 이상의 가중치 텐서는 CP 분해를 사용한다.

 

 

SVD

파이토치에서는 svd 함수를 제공한다.

M = torch.rand((4, 3))
U, s, V = torch.svd(M)
print(U)
print(s)
print(V)
tensor([[-0.7113, -0.6554, -0.0044],
        [-0.1999, -0.0309,  0.7525],
        [-0.1900, -0.0704, -0.6577],
        [-0.6465,  0.7513, -0.0346]])
tensor([1.2863, 0.4002, 0.1417])
tensor([[-0.4716,  0.0485,  0.8805],
        [-0.7004, -0.6272, -0.3406],
        [-0.5357,  0.7773, -0.3298]])

 

 

svd로 분해한 결과 중 U, V는 직교 행렬이다. (부동 소수점으로 인해 약간의 오차는 있음)

# 직교 행렬인지 확인
is_orthogonal_U = torch.allclose(torch.mm(U.t(), U), torch.eye(U.shape[1]), atol=1e-5)  # U^T * U == I 여야 함
is_orthogonal_V = torch.allclose(torch.mm(V.t(), V), torch.eye(V.shape[1]), atol=1e-5)  # V^T * V == I 여야 함

print("U는 직교 행렬인가요?:", is_orthogonal_U)
print("V는 직교 행렬인가요?:", is_orthogonal_V)
U는 직교 행렬인가요?: True
V는 직교 행렬인가요?: True

 

 

분해된 행렬을 다시 원래 행렬로 다시 결합할 수 있다. torch.mm은 행렬곱을 수행하고, torch.diag는 해당 벡터를 대각 성분으로 가지는 대각 행렬을 반환한다.

composed_M = torch.mm(torch.mm(U, torch.diag(s)), V.t())
print(M)
print(composed_M)
tensor([[0.4182, 0.8055, 0.2865],
        [0.2146, 0.1516, 0.0930],
        [0.0318, 0.2206, 0.1397],
        [0.4024, 0.3956, 0.6808]])
tensor([[0.4182, 0.8055, 0.2865],
        [0.2146, 0.1516, 0.0930],
        [0.0318, 0.2206, 0.1397],
        [0.4024, 0.3956, 0.6808]])

 

 

SVD low rank decomposition도 제공한다.

M = torch.rand((4, 3))
k = 2

Uk, sk, Vk = torch.svd_lowrank(M, q=k)
approximated_M = torch.mm(torch.mm(Uk, torch.diag(sk)), Vk.t())

print(M)
print(approximated_M)
tensor([[0.6078, 0.8166, 0.2079],
        [0.6972, 0.8622, 0.5827],
        [0.2672, 0.8073, 0.0417],
        [0.8150, 0.4247, 0.9876]])
tensor([[0.5244, 0.8510, 0.2628],
        [0.7237, 0.8513, 0.5652],
        [0.3254, 0.7833, 0.0033],
        [0.8218, 0.4219, 0.9831]])

 

 

CP(Canonical Polyadic) 분해

텐서를 여러 개의 요소의 외적(outer product)으로 분해하는 기법이다. 일반적으로 SVD보다 더 유연하고 효율적이다. CP 분해는 특히 다차원 데이터에서 패턴을 찾거나 잠재적인 구조를 추출하는 데 사용된다.

 

CP 분해는 다음과 같은 세 가지 유형이 있다:

  1. PARAFAC(Parallel Factor Analysis): 가장 기본적인 CP 분해 방법으로, 텐서를 여러 개의 외적으로 분해합니다. 각 외적은 텐서의 각 차원을 따라서 진행된다.
  2. CANDECOMP/CANDECOMP-PARAFAC(CP): 이 방법은 PARAFAC의 일반화 버전으로, 텐서의 각 차원마다 서로 다른 요소들을 분해한다. 각 차원에 대해 여러 개의 외적을 수행하며, 이러한 과정은 텐서의 특정 차원에 대한 분해를 가능하게 한다.
  3. Tucker Decomposition: CP 분해의 확장된 버전으로, 추가로 각 차원에 대한 모드 수(Mode size)를 지정한다. 이는 PARAFAC 및 CP 분해에서 발생하는 문제를 해결하기 위해 사용된다.

예를 들어, 3차원 텐서 $X$가 주어졌을 때, CP 분해는 다음과 같이 표현될 수 있습니다:

 

$X = \sum_{r=1}^R A_r \otimes B_r \otimes C_r$

여기서, $A_r, B_r, C_r$은 각각 텐서의 첫 번째, 두 번째, 세 번째 차원에 대한 벡터들의 집합이며, $\otimes$는 외적(outer product)을 나타낸다. $은 분해의 순위(rank)이며, 이는 분해된 요소의 개수를 의미한다.

 

CP 분해에서 사용되는 rank는 분해된 텐서에서 각 요소(벡터나 매트릭스)의 개수를 결정하는 매개변수이다. 이를 간단히 설명하면, rank는 분해된 텐서가 원래 텐서를 얼마나 잘 근사할 수 있는지를 나타내는 지표로 이해할 수 있다.

 

보통 CP 분해에서는 rank를 작은 값으로 설정하여 텐서를 더 낮은 차원으로 효과적으로 압축하고자 한다. 더 작은 rank를 사용하면 분해된 텐서가 원래 텐서를 더 간결하게 나타낼 수 있지만, 동시에 원래 텐서를 완벽하게 복원하는 데 필요한 정보의 손실이 발생할 수 있다.

 

일반적으로, 높은 rank를 사용하면 분해된 텐서가 원래 텐서를 더 정확하게 근사할 수 있지만, 이는 더 많은 메모리 및 계산 비용이 필요하게 됩니다. 따라서, rank를 선택하는 것은 성능과 비용 사이의 균형을 고려하여야 한다.

 

CP 분해에서의 rank는 각각의 요소(벡터나 매트릭스)의 차원이며, 이 값이 작을수록 분해된 텐서는 원래 텐서를 더 간결하게 나타낸다. 따라서, rank를 적절히 선택하는 것이 중요하다. 일반적으로, rank는 분해하려는 텐서의 특성과 분해 후의 용도에 따라 조정된다.

 

이렇게 CP 분해를 통해 텐서를 분해하면 원래 텐서를 해당 요소들의 외적으로 다시 조합하여 근사할 수 있다.

 

 

다음은 TensorLy 라이브러리로 CP 분해를 실습해 본다.

 

TensorLy의 백엔드를 파이토치로 설정한다. 이렇게하면 TensorLy는 이후의 연산을 파이토치를 통해 수행한다.

import torch
import tensorly as tl
from torch import nn
from torchvision import models
from tensorly import decomposition

tl.set_backend('pytorch')

 

 

우리가 경량화하고자하는 계층은 Conv2d로 각 차원이 의미하는 것은 [out_dim, in_dim, kernel_height, kernel_width] 이다. 

model = models.vgg16(weights="VGG16_Weights.IMAGENET1K_V1")

layer = model.features[0]
print(layer)
print("\n", layer.weight.data.shape)
Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))

 torch.Size([64, 3, 3, 3])

 

 

decomposition의 parafac 함수를 통해 텐서를 분해한다. weights는 [16] 크기의 벡터로 분해된 텐서의 중요도를 의미하는 가중치를 담고있다. 여기서는 normalize_factors=False로 지정하여 정규화되지 않으므로 값이 모두 1이다. factors에는 각 차원에 대해 분해한 텐서를 리스트에 담고 있다.

rank=16

weights, factors = decomposition.parafac(
    tensor=layer.weight.data,
    rank=rank,
    init="random",
    normalize_factors=False,
)

last, first, vertical, horizontal = factors

print(last.shape)
print(first.shape)
print(vertical.shape)
print(horizontal.shape)
torch.Size([64, 16])
torch.Size([3, 16])
torch.Size([3, 16])
torch.Size([3, 16])

 

 

이렇게 분해된 합성곱 계층의 가중치를 활용해 새로운 합성곱 계층을 정의한다. 아래 코드는 합성곱 계층을 CP 분해를 통해 네 개의 합성곱 계층으로 분해하는 함수를 정의한다. 각 합성곱 계층이 의미하는 것은 다음과 같다.

  1. pointwise_s_to_r_layer:
    • 이 레이어는 분해된 첫 번째 요소인 first를 입력으로 받고, 분해된 마지막 요소인 last를 출력으로 생성합니다.
    • 입력 텐서의 채널 수를 출력 텐서의 채널 수로 변환합니다.
    • 커널 크기가 1x1인 합성곱 레이어이므로, 이를 pointwise 합성곱이라고 합니다.
    • 즉, 입력 텐서의 각 픽셀에 대해 각 채널에 대한 선형 변환을 수행합니다.
  2. depthwise_vertical_layer:
    • 이 레이어는 분해된 vertical 요소를 입력으로 받고, 동일한 채널 간에만 합성곱을 적용합니다.
    • 입력 텐서의 채널 간에만 합성곱을 적용하므로, Depthwise Separable Convolution이라고도 합니다.
    • 입력 텐서의 높이 방향으로 1x1 커널을 적용하여 세로 방향의 특징을 추출합니다.
    • 따라서, 입력 텐서의 채널 수는 유지되고, 높이 방향의 크기는 줄어듭니다.
  3. depthwise_horizontal_layer:
    • 이 레이어는 분해된 horizontal 요소를 입력으로 받고, 동일한 채널 간에만 합성곱을 적용합니다.
    • 입력 텐서의 채널 간에만 합성곱을 적용하므로, Depthwise Separable Convolution이라고도 합니다.
    • 입력 텐서의 너비 방향으로 1x1 커널을 적용하여 가로 방향의 특징을 추출합니다.
    • 따라서, 입력 텐서의 채널 수는 유지되고, 너비 방향의 크기는 줄어듭니다.
  4. pointwise_r_to_t_layer:
    • 이 레이어는 분해된 last 요소를 입력으로 받고, 원래 출력 채널 수로 변환합니다.
    • 입력 텐서의 채널 수를 출력 텐서의 채널 수로 변환합니다.
    • 커널 크기가 1x1인 합성곱 레이어이므로, 이를 pointwise 합성곱이라고 합니다.
    • 즉, 입력 텐서의 각 픽셀에 대해 각 채널에 대한 선형 변환을 수행합니다.
def cp_decomposition(layer, rank):
  weights, factors = decomposition.parafac(
      tensor=layer.weight.data,
      rank=rank,
      init="random",
      normalize_factors=False
  )
  last, first, vertical, horizontal = factors

  pointwise_s_to_r_layer = nn.Conv2d(
      first.shape[0],
      first.shape[1],
      kernel_size=1,
      stride=1,
      padding=0,
      dilation=layer.dilation,
      bias=False,
  )
  depthwise_vertical_layer = nn.Conv2d(
      vertical.shape[1],
      vertical.shape[1],
      kernel_size=(vertical.shape[0], 1),
      stride=1,
      padding=(layer.padding[0], 0),
      dilation=layer.dilation,
      groups=vertical.shape[1],
      bias=False,
  )
  depthwise_horizontal_layer = nn.Conv2d(
      horizontal.shape[1],
      horizontal.shape[1],
      kernel_size=(1, horizontal.shape[0]),
      stride=layer.stride,
      padding=(0, layer.padding[0]),
      dilation=layer.dilation,
      groups=horizontal.shape[1],
      bias=False,
  )
  pointwise_r_to_t_layer = nn.Conv2d(
      last.shape[1],
      last.shape[0],
      kernel_size=1,
      stride=1,
      padding=0,
      dilation=layer.dilation,
      bias=True,
  )
  pointwise_r_to_t_layer.bias.data = layer.bias.data

  depthwise_horizontal_layer.weight.data = (
      torch.transpose(horizontal, 1, 0).unsqueeze(1).unsqueeze(1)
  )
  depthwise_vertical_layer.weight.data = (
      torch.transpose(vertical, 1, 0).unsqueeze(1).unsqueeze(-1)
  )
  pointwise_s_to_r_layer.weight.data = (
      torch.transpose(first, 1, 0).unsqueeze(-1).unsqueeze(-1)
  )
  pointwise_r_to_t_layer.weight.data = last.unsqueeze(-1).unsqueeze(-1)

  new_layers = [
      pointwise_s_to_r_layer,
      depthwise_vertical_layer,
      depthwise_horizontal_layer,
      pointwise_r_to_t_layer,
  ]
  return nn.Sequential(*new_layers)

 

 

CP 분해를 수행하면 가중치 수가 줄어드는 것을 확인할 수 있다.

layer_cp_decomposed = cp_decomposition(layer, rank=16)

print("CP 분해 전 가중치 수:", sum(param.numel() for param in layer.parameters()))
print("CP 분해 후 가중치 수:", sum(param.numel() for param in layer_cp_decomposed.parameters()))
CP 분해 전 가중치 수: 1792
CP 분해 후 가중치 수: 1232

 

 

하나의 합성곱 계층만 분해하는 예시를 보았지만, 이 함수를 이용해 다음과 같이 VGG16의 모든 합성곱 계층을 분해하는 것이 가능하다.

import copy


decomposed_model = copy.deepcopy(model)
for idx, module in enumerate(decomposed_model.features):
    if isinstance(module, nn.Conv2d):
        rank = max(module.weight.data.numpy().shape) // 3
        decomposed_model.features[idx] = cp_decomposition(module, rank)

print("CP 분해 전 가중치 수 :", sum(param.numel() for param in model.parameters()))
print("CP 분해 후 가중치 수 :", sum(param.numel() for param in decomposed_model.parameters()))
CP 분해 전 가중치 수 : 138357544
CP 분해 후 가중치 수 : 124799037

 

 

3. ONNX

ONNX(Open Neural Network Exchange)는 딥러닝 모델의 구조와 가중치를 표현하기 위한 중립적인 형식을 제공하여 프레임워크 간에 호환성을 가지는 오픈소스 플랫폼이다. 다양한 딥러닝 프레임워크 간에 모델을 변환하고 공유할 수 있는 기능을 제공하는 중간 언어라고 할 수 있다.

 

ONNX를 지원하는 프레임워크는 ONNX 형식으로 저장된 모델을 불러와 실행할 수 있으며, ONNX를 통해 학습된 모델을 다른 프레임워크로 이식할 수도 있다.

 

ONNX 런타임은 ONNX를 실행하기 위한 오픈소스 엔진으로, ONNX 모델을 부럴와 실행하는 역할을 담당하며, ONNX 형식으로 저장된 모델을 효율적으로 실행하기 위해서 최적화된 엔진을 제공한다. ONNX 런타임은 파이토치, 텐서플로우 등의 주요 딥러닝 프레임워크와의 연동을 지원한다. 

 

이를 통해 고성능의 추론 기능을 제공하면서 가벼운 런타임 환경을 구성할 수 있으며,  모바일 기기와 같은 리소스가 제한된 환경에서도 효율적으로 ONNX 모델을 실행할 수 있다.

 

 

import urllib
import time
import torch
import onnxruntime as ort
from PIL import Image
from torch import onnx
from torchvision import models, transforms

 

ONNX 형식 변환

onnx.export 함수로 파이토치 모델을 ONNX 형식으로 내보낼 수 있다. args는 모델에 전달할 입력, f는 저장할 경로를 의미한다.

model = models.vgg16(weights="DEFAULT")
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)
onnx.export(
    model=model,
    args=dummy_input,
    f="vgg16.onnx",
)

 

 

ONNX 런타임 실행

파이토치와 onnx를 비교해보기 위해 우선 파이토치를 사용했을 때의 출력과 추론 시간을 확인한다.

model = models.vgg16(weights="DEFAULT")
model.eval()

with torch.no_grad():
    start_time = time.time()
    output = model(input)
    end_time = time.time()
    print("파이토치:")
    print(output[:, :10])
    print(end_time - start_time)
파이토치:
tensor([[-47.1957,  17.0475,   0.9968, -54.8332,  -5.2323, -13.0191, -16.3137,
          14.8506,  33.1292, -47.8647]])
0.31764888763427734

 

 

  • ort.InferenceSession 클래스로 ONNX 모델을 로드하고 세션을 할당한다.
  • ort_session.get_inputs 메서드는 ONNX 모델에서 정의된 입력에 대한 정보를 가져온다. 이를 통해 ONNX 모델 추론에 필요한 입력을 생성한다. 
  • ort_session.run 메서드는 ONNX 모델로 추론을 수행한다.
def to_numpy(tensor):
  return tensor.detach().cpu().numpy()
  
ort_session = ort.InferenceSession("vgg16.onnx")

start_time = time.time()
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(input)}
ort_outs = ort_session.run(output_names=None, input_feed=ort_inputs)
end_time = time.time()
print("ONNX:")
print(ort_outs[0][:, :10])
print(end_time - start_time)
ONNX:
[[-47.19573     17.04755      0.99684036 -54.833183    -5.23222
  -13.0191     -16.313707    14.850587    33.12925    -47.86475   ]]
0.14706873893737793

 

결과를 보면 추론시간이 감소하였다.