본문 바로가기

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

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

모델 경량화(Model Compression)란 모델 크기와 연산량을 줄여서 작은 메모리 공간과 적은 연산량으로도 효율적으로 작동할 수 있게 만드는 프로세스를 의미한다. 

 

모바일이나 임베디드 환경에서는 제한된 메모리 및 연산 능력을 가지고 있으므로 모델 경량화가 중요하다. 임베디드 환경(Embedded Environment)은 늑정한 기능이나 작업을 수행하기 위해 설계된 전용 시스템이나 장치를 가리킨다. 이러한 임베디드 시스템은 주로 특정한 작업에 중점을 둔 하드웨어와 소프트웨어의 조합으로 이루어져 있다. 예를 들면 웨어러블 기기, 자율 주행 차량 등이 있다.

 

모델 경량화의 주요 이점은 다음과 같다.

  • 모델 크기 감소 : 작은 모델은 메모리 요구사항이 적어 메모리 용량 제약이 있는 모바일 장치나 임베디드 시스템에서도 용이하게 실행될 수 있다. 또한 모델을 클라우드로 전송하는 데 걸리는 시간과 비용을 감소시킬 수 있다.
  • 연산량 최적화 : 모델 내부의 연산량을 최적화하여 더 적은 연산 리소스를 사용하도록 만든다. 모바일 장치나 임베디드 시스템의 배터리 수명을 연장하고 발열을 감소시킬 수 있으며, 클라우드 환경에서 실행 비용을 절감하는 데 도움을 준다.
  • 예측 속도 향상 : 모바일 애플리케이션에서는 실시간 응답이 필요한 경우가 많은데, 작은 모델은 더 빠른 예측을 가능하게 하여 사용자 경험을 향상할 수 있다.

 

 

1. Pruning

Pruning은 모델의 매개변수 중 중요도가 낮은 매개변수를 제거하여 전체 매개변수 수를 줄이는 방법이다. Pruning을 수행한 후에는 fine-tuning이 필요하다.

 

 

1.1 Unstructured Pruning

모델의 매개변수 중 일부를 0으로 만드는 방법이다. 

 

Magnitude-based Pruning

가중치의 크기에 기반하여 가중치를 선택적으로 제거하는 방법이다. 주로 L1 Norm 또는 L2 Norm이 가중치의 크기를 측정하는 데 사용된다. 일반적으로 미리 정의된 임계값보다 작은 크기의 가중치는 제거된다. 작은 가중치는 모델의 성능에 미치는 영향이 적다고 가정한다.

 

Gradient-based Pruning

모델의 학습 중에 gradient를 기반으로 가중치를 제거하는 방법이다. 학습 중에 가중치의 gradient가 작은 경우 해당 가중치는 중요하지 않다고 판단하여 제거한다.

 

 

다음은 BERT 모델에 unstructured pruning을 수행하는 예시이다.

  • torch.nn.utils 패키지의 prune 모듈은 pruning에 사용되는 다양한 함수를 제공한다.
  • global_unstructured 함수의 parameters에 pruning을 적용한 레이어들을 리스트로 전달한다. 여기서는 L1 Norm을 사용하였고, amount=0.2로 설정하여 20%의 매개변수를 제거한다.
import torch
from torch.nn.utils import prune
from transformers import BertTokenizer, BertForSequenceClassification

tokenizer = BertTokenizer.from_pretrained(
    pretrained_model_name_or_path="bert-base-multilingual-cased",
    do_lower_case=False,
)
model = BertForSequenceClassification.from_pretrained(
    pretrained_model_name_or_path="bert-base-multilingual-cased",
    num_labels=2
)

print("가지치기 적용 전:")
print(model.bert.encoder.layer[0].attention.self.key.weight)

parameters = [
    (model.bert.embeddings.word_embeddings, "weight"),
    (model.bert.encoder.layer[0].attention.self.key, "weight"),
    (model.bert.encoder.layer[1].attention.self.key, "weight"),
    (model.bert.encoder.layer[2].attention.self.key, "weight"),
]

prune.global_unstructured(
    parameters,
    pruning_method=prune.L1Unstructured,
    amount=0.2,
)

print("가지치기 적용 후:")
print(model.bert.encoder.layer[0].attention.self.key.weight)
가지치기 적용 전:
Parameter containing:
tensor([[-0.0259, -0.0599, -0.0395,  ...,  0.0083, -0.0271, -0.0043],
        [ 0.0419,  0.0561,  0.0044,  ..., -0.0042,  0.0200,  0.0163],
        [ 0.0160, -0.0439, -0.0092,  ..., -0.0115, -0.0879, -0.0514],
        ...,
        [-0.0310, -0.0218,  0.0345,  ..., -0.0010,  0.0107,  0.0088],
        [-0.0756, -0.0755, -0.0067,  ..., -0.0132, -0.0330, -0.0986],
        [ 0.0159, -0.0423, -0.0061,  ...,  0.0445, -0.0184,  0.0410]],
       requires_grad=True)

가지치기 적용 후:
tensor([[-0.0259, -0.0599, -0.0395,  ...,  0.0000, -0.0271, -0.0000],
        [ 0.0419,  0.0561,  0.0000,  ..., -0.0000,  0.0200,  0.0163],
        [ 0.0160, -0.0439, -0.0000,  ..., -0.0115, -0.0879, -0.0514],
        ...,
        [-0.0310, -0.0218,  0.0345,  ..., -0.0000,  0.0000,  0.0000],
        [-0.0756, -0.0755, -0.0000,  ..., -0.0132, -0.0330, -0.0986],
        [ 0.0159, -0.0423, -0.0000,  ...,  0.0445, -0.0184,  0.0410]],
       grad_fn=<MulBackward0>)

 

결과를 보면 일부 값이 작은 가중치들이 0.0으로 변경됐다.

 

 

1.2 Structured Pruning

Structured Pruning은 딥러닝 모델에서 특정한 구조를 가진 가중치 블록을 동시에 제거하거나 가중치 값을 감소시키는 기술이다. Structured Pruning은 Unstructured Pruning과 달리, 모델 내에서 일부 특정한 패턴이나 블록을 선택하여 제거한다. 이는 특정한 가중치가 아닌 가중치 블록을 삭제하므로 모델이 가지는 구조를 유지하면서도 모델 크기를 줄일 수 있다.

 

 

Filter Pruning

CNN에서 주로 사용되며, 특정 필터를 삭제하여 모델의 크기를 줄이는 방식이다. 작은 가중치를 가진 필터나 중요하지 않은 필터를 삭제하여 유용한 특징만을 남긴다.

 

Channel Pruning

더 큰 구조적인 단위로, 특정 채널을 삭제하여 모델 크기를 줄인다. 각 채널은 특정 유형의 특징을 감지하므로, 중요하지 않거나 중복되는 채널을 제거함으로써 모델을 간소화할 수 있다.

 

Structured Sparsity Inducing Regularization

특정한 가중치의 블록에 대해 L1 또는 L2 규제를 적용하여, 해당 블록의 가중치를 감소시키거나 제거하는 방식입니다.

 

 

2. Quantization

인공 신경망 모델이 32비트 형식으로 저장되어있으면 메모리 내 공간을 많이 차지한다. 양자화(Quantization)는 모델의 가중치를 제한된 비트 수로 표현하여 모델의 메모리 사용량을 줄이고 추론 속도를 향상시킨다.

 

정적 양자화(Static Quantization)

양자화 매개변수를 미리 계산하여 고정된 값으로 사용하는 방법이다. 이를 위해 데이터셋의 일부를 입력해보고, 양자화 매개변수를 계산하는 교정(calibration) 과정이 필요하다. 정적 양자화는 고정된 매개변수를 사용하기 때문에 추론 시 추가적인 연산이 필요하지 않다.

 

동적 양자화(Dynamic Quantization)

모델의 가중치는 미리 양자화하지만, 활성화 함수는 데이터가 입력될 때마다 양자화 매개변수를 계산한다. 이로 인해 정적 양자화보다는 더 많은 계산 비용이 들지만, 입력값에 맞춰 매개변수가 연산되기 때문에 정적 양자화보다 더 나은 성능을 보일 수 있다.

 

이 외에도 양자화를 진행하는 시점에 따라 학습 후 양자화(Post Training Quantization)양자화 인식 학습(Quantization Aware Training)으로 구분할 수 있다.

 

 

2.1 Post Training Quantization

학습 후 양자화는 모델을 학습한 후에 양자화를 적용하는 방법으로 이미 학습된 모델을 양자화하여 추론 시에 사용된다.

 

 

VGG16 학습 후 정적 양자화 코드의 일부

QuantizedVGG16 클래스는 사전 학습된 VGG16 모델의 입력부와 출력부에 양자화 스텁(QuantStub)비양자화 스텁(DeQuantStub)을 연결한다. 양자화 스텁은 모델의 입력에 대한 양자화 연산을 수행하고 비양자화 스텁은 양자화된 모델의 출력에 대한 비양자화 연산을 수행한다.

class QuantizedVGG16(nn.Module):
  def __init__(self, model_fp32):
    super(QuantizedVGG16, self).__init__()
    self.quant = quantization.QuantStub()
    self.dequant = quantization.DeQuantStub()
    self.model_fp32 = model_fp32

  def forward(self, x):
    x = self.quant(x)
    x = self.model_fp32(x)
    x = self.dequant(x)
    return x

 

 

양자화를 수행하기 위해서는 양자화를 적용하려는 모델의 qconfig 속성에 양자화 백엔드를 할당해야 한다. 그런 다음 quantization의 prepare 함수로 모델의 가중치를 양자화 가능한 형태로 변환한다.

import torch
from torch.ao import quantization
from torchvision import models


model = models.vgg16(num_classes=2)
model.load_state_dict(torch.load(ROOT_PATH / "models" / "VGG16.pt"))

quantized_model = QuantizedVGG16(model).to(device)

quantization_backend = "fbgemm"
quantized_model.qconfig = quantization.get_default_qconfig(quantization_backend)

model_static_quantized = quantization.prepare(quantized_model)

 

 

양자화된 모델로 변경 됐다면 양자화 매개변수를 계산하는 교정 과정을 수행한다.

for i, (image, target) in enumerate(calibartion_dataloader):
  if i >= 10:
    break
  model_static_quantized(image.to(device))

 

 

교정이 완료되면 양자화 모델로 변환시키기 위해 CPU로 변환한다. 현재 파이토치의 양자화는 CPU만 지원한다. 이후 convert 함수로 모델 내의 양자화 관련 연산자들을 실제 양자화 연산자로 대체한다.

model_static_quantized.to(torch.device("cpu"))
model_static_quantized = quantization.convert(model_static_quantized)

 

 

모델이 양자화됐다면 torch.jit.script 함수를 통해 스크립트 모듈로 변환한다. 스크립트 모듈은 그래프 형태로 저장되어 모델의 실행 그래프를 최적화하고 추론 시에 불필요한 오버헤드를 줄일 수 있다. 스크립트화된 모델은 torch.jit.save 함수로 저장한다.

torch.jit.save(torch.jit.script(model_static_quantized), ROOT_PATH / "models" / "PTSQ_VGG16.pt")

 

 

다음은 양자화 결과 비교이다.

with torch.no_grad():
  start_time = time.time()
  outputs = model(inputs)
  file_size = os.path.getsize(ROOT_PATH / "models" / "VGG16.pt") / 1e6
  print("양자화 적용 전:")
  print(f"추론 시간: {time.time() - start_time:.4f}s")
  print(f"파일 크기: {file_size:.2f} MB")
  print('\n')

start_time = time.time()
outputs = model_static_quantized(inputs)
file_size = os.path.getsize(ROOT_PATH / "models" / "PTSQ_VGG16.pt") / 1e6

print("양자화 적용 후:")
print(f"추론 시간: {time.time() - start_time:.4f}s")
print(f"파일 크기: {file_size:.2f} MB")
양자화 적용 전:
추론 시간: 0.1266s
파일 크기: 537.08 MB


양자화 적용 후:
추론 시간: 0.0942s
파일 크기: 134.54 MB

 

 

VGG16 학습 후 동적 양자화

동적 양자화는 모델을 실행하는 동안에만 필요한 부분을 양자화하는 기법이다. 모델을 불러오고 평가모드로 전환한 뒤 quantize_dynamic 함수를 통해 수행한다. model에 양자화 하려는 모델을 전달, qconfig_spec에는 양자화하려는 계층, dtype에는 양자화 데이터 형식을 전달한다.

model = models.vgg16(num_classes=2)
model.load_state_dict(torch.load(ROOT_PATH / "models" / "VGG16.pt"))
model = model.to(device)
model.eval()

model_dynamic_quantized = quantization.quantize_dynamic(
    model=model,
    qconfig_spec={nn.Linear},
    dtype=torch.qint8,
)
model_dynamic_quantized.eval()

torch.save(model_dynamic_quantized.state_dict(), ROOT_PATH / "models" / "PTDQ_VGG16.pt")

 

 

양자화를 수행하고나면 Linear 계층이 DynamicQuantizedLinear 계층으로 변환된다.

file_size = os.path.getsize(ROOT_PATH / "models" / "VGG16.pt") / 1e6
print("양자화 적용 전")
print(f"파일 크기: {file_size:.2f} MB")
print(model.classifier)
print("\n")

file_size = os.path.getsize(ROOT_PATH / "models" / "PTDQ_VGG16.pt") / 1e6
print("양자화 적용 후")
print(f"파일 크기: {file_size:.2f} MB")
print(model_dynamic_quantized.classifier)
양자화 적용 전
파일 크기: 537.08 MB
Sequential(
  (0): Linear(in_features=25088, out_features=4096, bias=True)
  (1): ReLU(inplace=True)
  (2): Dropout(p=0.5, inplace=False)
  (3): Linear(in_features=4096, out_features=4096, bias=True)
  (4): ReLU(inplace=True)
  (5): Dropout(p=0.5, inplace=False)
  (6): Linear(in_features=4096, out_features=2, bias=True)
)


양자화 적용 후
파일 크기: 178.45 MB
Sequential(
  (0): DynamicQuantizedLinear(in_features=25088, out_features=4096, dtype=torch.qint8, qscheme=torch.per_tensor_affine)
  (1): ReLU(inplace=True)
  (2): Dropout(p=0.5, inplace=False)
  (3): DynamicQuantizedLinear(in_features=4096, out_features=4096, dtype=torch.qint8, qscheme=torch.per_tensor_affine)
  (4): ReLU(inplace=True)
  (5): Dropout(p=0.5, inplace=False)
  (6): DynamicQuantizedLinear(in_features=4096, out_features=2, dtype=torch.qint8, qscheme=torch.per_tensor_affine)
)

 

 

2.2 Quantization Aware Training, QAT

양자화 인식 학습은 모델의 학습 시점에 양자화로 인한 영향을 미리 계산하는 방법이다. 신경망을 학습하는 도중, 가짜 양자화 모듈을 신경망에 배치하여 신경망이 양자화되었을 때의 효과와 양자화 매개변수를 계산한다. 학습이 끝나면 가짜 양자화 모듈에 저장된 정보를 이용하여 신경망을 양자화한다. 학습 후 추가적인 학습이 필요하지만, 양자화 후의 성능 저하가 가장 적다.

 

 

일반적인 모델 학습 과정과 정적 양자화를 병합하여 수행하는데, prepare 함수를 사용하는 것이 아닌 prepare_qat 함수를 사용한다. 이후 모델을 학습하고 정적 양자화와 마찬가지로 torch.jit.script를 통해 스크립트 모듈로 변환하여 저장한다.

model = models.vgg16(weights="VGG16_Weights.IMAGENET1K_V1")
for param in model.parameters():
    param.requires_grad = False
model.classifier[6] = nn.Linear(4096, len(train_dataset.classes))

quantization_backend = "fbgemm"
quantized_model = QuantizedVGG16(model).to(device)
quantized_model.qconfig = quantization.get_default_qat_qconfig(quantization_backend)
quantization.prepare_qat(quantized_model)