[Torch] 딥러닝 모델 경량화 방법 정리(Pruning, Quantization, 모바일 배포 방법, ONNX)

- 이번 포스팅에서는 딥러닝 모델 경량화 방법을 정리하겠습니다.

최근 저는 개인적인 어플리케이션 개발과 더불어 실험적인 구조의 Image Segmentation 모델을 학습시키는 중인데,

제가 만든 Image Segmentation 모델을 모바일/드론 등 엣지 디바이스에 적용하여 실시간(30FPS)으로 시각 정보를 처리하는 것을 목표로 하고 있습니다.

모델 학습 경과

 

이렇게 모바일 환경과 같은 제한적인 환경에서 비교적 무거운 딥러닝 모델을 구동시키기 위해 필요한 몇가지 기법들을 정리하겠습니다.

 

(Pruning(가지치기))

- Pruning 이란, 가지치기라는 뜻과 같이 모델 구조에서 중요도가 떨어지는 부분을 잘라내는 기법을 의미합니다.

 

- 쉽게 설명하자면,

f(x) = x * w + b

라는 수식이 무수히 많이 붙은 레이어에서, w 값이 극도로 낮은 수식은 모델에서 제거를 하는 개념입니다.

아시다시피 w 가 극단적으로 0에 가깝다고 한다면, 해당 수식이 결과값에 끼치는 영향을 극도로 낮기 때문에,

쓸데없이 파라미터로써 메모리를 차지하고 순전파시 연산 비용을 늘리는 이러한 부분을 제거함으로써 모델 전체적인 사이즈를 줄일 수 있습니다.

 

- Pruning 의 장점은,

위에서 설명했듯, 모델 퀄리티에 영향을 끼치지 않는 부분을 잘라냄으로써 모델의 경량화를 이룰 수 있다는 것이고,

단점은, 그렇게 작은 영향이라도 중첩되었을 경우 무시 못할 영향력을 가지는 구조인 경우 정확도에 악영향을 끼칠 수 있다는 것입니다.

이러한 단점으로 인해 Pruning 이후에는 Fine-Tuning 을 수행하여 파라미터의 학습을 재정렬 해주어야만 합니다.

 

- 위와 같은 Pruning 의 장점과 특징을 이용한 딥러닝 모델 설계 방식을 알려드리겠습니다.

 

일반적으로 딥러닝은 모델의 크기가 크고 깊이가 깊을수록 더 나은 판단을 한다는 것이 증명되었습니다.

기본적인 딥 레이어 퍼셉트론 구조는 물론이고, 보다 목적과 특징에 특화된 Convolutional 네트워크나, RNN, Transformer Block 등 검증된 여러 정형화된 구조들이 존재하지만, 어느 아키텍쳐와 블록을 사용하더라도, 모델의 파라미터가 많고 깊이가 깊을수록 성능이 높아지는 것은 불변의 법칙이죠.

 

딥러닝 관련 본인의 설계를 해보신 분이시라면 아시겠지만,

적절한 파라미터를 찾는 공식 같은 것은 없습니다.

 

모델 구조에 대한 아이디어가 있고, 이를 실현시키려고 해도 각 파트별 가중치 사이즈를 어떻게 할 지에 대해서는 현 시점에서는 순전히 운에 맡겨야 하는 부분입니다.

 

되도록 파라미터를 많이 하고 싶을 수도 있는데, 전체 구조에서 일부 파라미터만 높여봐야 다음 단계에서 파라미터가 작다면 대부분의 정보가 잘려나가 낭비될 것이기에 일반적으로 한 부분을 증량하면 전체를 증량해야 하는 것은 당연한 일입니다.

(다음 출력 벡터의 의미조차 우리가 전부 정의내릴 수 없기에, 정보를 압축한다는 의미에서 축약을 할 수도 있고, 정보의 조합으로 새로운 정보를 탐색한다는 의미에서 벡터를 확장할 수도 있으므로 축약할지 확장할지에 대한 기준도 없습니다.)

 

문제는, 딥러닝 모델의 좋고 나쁨은 이미 검증된 모델을 사용하지 않는 하에는 직접 학습을 진행하고 그 결과를 확인해야만 한다는 것입니다.

 

이때, 적절한 값을 찾기 위하여 작은 파라미터를 사용하기보다는,

하드웨어가 허용하는 제법 큰 파라미터를 처음부터 사용하여 성능을 확인하고,

만약 만족할 만한 성능이 나온다면 이를 Pruning 을 하여 최적의 사이즈를 찾아나가는 방식을 사용할 수도 있을 것입니다.

 

이러한 방식은 실제로 딥러닝 연구/설계자들 사이에서도 자주 사용되는 기법이라고 합니다.

 

- Pruning 을 적용하기 위한 Torch 코드는 아래와 같습니다.

import torch
import torch.nn.utils.prune as prune
import torch.nn as nn

# 1. 모델 클래스 정의
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, 3)
        self.conv2 = nn.Conv2d(6, 12, 3)
        self.fc = nn.Linear(12 * 6 * 6, 10)  # 예시용

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x

# 2. 모델 로드
model = MyModel()
model.load_state_dict(torch.load("model.pth", map_location="cpu"))

# 3. Pruning 적용 (예: conv1에 50% 가지치기)
prune.l1_unstructured(model.conv1, name="weight", amount=0.5)

# 원한다면 pruning mask 제거 후 진짜로 가지치기
prune.remove(model.conv1, "weight")  # pruning mask 제거 및 적용

# 4. 다시 저장
torch.save(model.state_dict(), "model_pruned.pth")

 

이미 위와 같은 모델을 학습시켜 나온 pth 파일이 있을 때,

이를 load_state_dict 로 불러온 후,

각 위치별 prune.unstructured 를 적용하여 가지치기를 할 수 있습니다.

 

l1_unstructured 는,

L1 Norm 을 기준으로 작은 weight 순서로 50% 를 제거(weight 를 기준으로 전체에서 50% 제거)하는 코드로,

이때에는 실제로 구조가 잘려나가는 것이 아니라, 해당 부분에 mask 처리를 하여 w 를 완전히 0으로 침묵시키는 것입니다.

 

prune.remove 를 해야만 masking 된 부분이 제거되며 모델이 경량화 되며,

이를 다시 저장하여, 보다 가벼워진 모델을 사용할 수 있게 됩니다.

 

(Quantization)

- 양자화란, 딥러닝 파라미터의 역할을 수행하는 실수 변수값의 표현 범위를 제한하는 방식입니다.

컴퓨터 사이언스 지식이 있으시다면 컴퓨터에서 데이터는 bit 의 조합으로 이루어진다는 것을 아실 것입니다.

실수의 경우 역시 bit 들의 조합으로 이루어지며, bit 의 개수가 많을수록 정밀도가 높습니다.

일반적인 딥러닝 파라미터의 실수 표현은 32bit 로 이루어집니다.

당연히 이 bit 크기가 커질수록 파라미터 하나를 계산할 때에 들어가는 연산량과 메모리 소모량이 많아지지만, 그만큼 계산 결과의 정확도와 표현력이 높아지고, 반면 bit 수가 작아진다면 메모리 소모량이 적어지고, 연산 속도도 빨라지지만 계산의 오차가 그만큼 쌓여서 결과의 정확도가 낮아질 수 있죠.

 

32bit 정밀도를 16bit, 8bit 로 줄여서 모델 사이즈를 줄이는 것을 모델 양자화라 부릅니다.

 

- 양자화 단위는 32bit, 16bit, 8bit, 4bit 까지 존재합니다.

심지어 실험적인 논문에서는 bit 수를 1bit 까지 낮추어 파라미터를 구성해도 전체 모델 성능이 크게 떨어지지 않았다는 결과가 있죠.

배포하려는 디바이스에 모델 사이즈를 맞추기 위해 가장 쉽게 고려할 수 있는 것이 바로 이 양자화입니다.

(단순 계산해서 32bit 짜리를 4bit 까지 양자화 했다고 하면, 모델 크기가 무려 8배나 줄어드는 것입니다.)

 

앞서 Pruning 과 마찬가지로 양자화 이후에는 간단히 파인튜닝을 해주는 것이 성능을 안정화 시킬 수 있습니다.

 

- 양자화를 하면 성능이 낮아질 수 있다는 것을 앞서 설명했는데,

이를 완화시킬 수 있는 방식이 있습니다.

아예 훈련 시점에 weights 에 양자화를 적용했다고 가정을 하여 표현범위를 제한하여 학습시키는 방식으로,

학습 시점과 양자화 시점간의 성능 차이의 간격을 줄이는 훈련 방식입니다.

이를 Quantization Aware Training(QAT) 라고 부릅니다.

 

(모바일 최적화)

- 서버와 같은 고정형 PC 환경과 달리 모바일과 같은 엣지 디바이스는 휴대성과 이동성이 있는 만큼 구동을 위하여 최대한 경량화가 필요합니다.

앞서 설명한 양자화나 Pruning, 모델 자체의 효율 상승을 제외하고 모델 배포 단계에서 고려할 것은, 어떤 기술을 이용하여 Torch 를 사용하여 만든 모델을 이식하느냐입니다.

 

- Tensorflow 를 사용한다면, 해당 모델을 바로 TensorflowLite 와 같은 경량 모바일용 추론 라이브러리로 배포해서 돌리는 방식이 있고,

Torch 를 사용한다면, Torch 모델을 TorchScript 모델로 변환 후, 각 모바일 OS 에서 지원해주는 라이브러리로 돌리는 방식이 있습니다.

 

torchscript 의 경우,

# PyTorch 모델 → TorchScript 모델로 변환
scripted_model = torch.jit.script(model)
scripted_model.save("model.pt")

 

위와 같이 torch 모델을 torch.jit.script 로 변환하여 이를 pt 파일로 저장하고,

이를 모바일로 옮겨, Android 에서는 pytorch_android 나 pythorch_android_lite 라이브러리로, iOS 에서는 LibTorch 나 TorchScript 로 추론을 할 수 있습니다.

 

각 모바일 전용 라이브러리는 모바일 최적화가 이루어져 최적의 성능 효율을 기대할 수 있습니다.

 

- Android 와 iOS 에만 국한되어 보자면,

NNAPI, CoreML 에 대해 설명드릴 수 있겠습니다.

 

NNAPI(DSP 사용) 는 Android 전용 추론 엔진, CoreML(Neural Engine 사용) 은 iOS 전용 추론 엔진으로,

각각 디바이스에 최적화된 딥러닝 추론 엔진으로, 당연히 이것을 이용해 딥러닝 모델을 실행시킨다면 최고 성능을 발휘할 수 있을 것입니다.

 

NNAPI 를 사용하려면, Tensorflow 의 경우는 TFLite 변환 후 이를 NNAPI 로 바로 사용하면 되고,

만약 Torch 를 사용한다면, Torch 모델을 ONNX 로 변환 후 이를 Tensorflow 변환 -> TFLite 변환 -> NNAPI 사용의 다소 복잡한 방식을 사용해야만 합니다.

 

다행해 CoreML 은 Tensorflow 뿐만 아니라 Torch 역시 바로 지원하므로 그냥 사용하면 됩니다.

 

(ONNX)

- 추가로, 딥러닝 배포에 관련된 기술 중 ONNX 에 대해 정리하겟습니다.

 

- ONNX(Open Neural Network Exchange)란,

딥러닝 모델 교환의 표준입니다.

딥러닝 라이브러리는 다양합니다.

크게는 Pytorch, Tensorflow 가 양립하고 있는데, 그외에도 몇가지 라이브러리가 있고 앞으로도 늘어날 가능성이 있죠.

라이브러리는 달라지지만, 딥러닝이라는 근본 구조는 변하지 않습니다.

 

이론적으론, 딥러닝 아키텍쳐 구조와 학습된 파라미터만 있다면 어떤 라이브러리로 학습시킨 모델이라도 상관 없이 실행시킬 수 있어야 하죠.

 

ONNX 라는 표준은 Pytorch, Tensorflow, MXNet 등의 다양한 딥러닝 프레임워크간 모델 호환을 가능하게 하는 일종의 표준어라고 이해하시면 됩니다.

 

torch 의 딥러닝 모델을 먼저 .onnx 파일로 변환한 후, 이를 배포에 사용하는 방식으로 사용이 가능합니다.

import torch.onnx

# PyTorch 모델을 ONNX로 변환
torch.onnx.export(model, dummy_input, "model.onnx", opset_version=11)

 

주의해야 할 점은,

아무리 플랫폼 독립적인 기술이라고 해도, 간혹 최신 버전 문법이나 if, for 과 같은 동적 제어 흐름을 해석할 때에는 에러가 날 가능성이 있다는 것을 주의하세요.

 

- 추론용 모델 배포 방식중 현 시점 가장 성능이 좋은 것은,

ONNX -> TensorRT 라고 알려져 있습니다.

 

TensorRT란,

NVIDIA 에서 제공하는 고성능 추론 최적화 및 실행 엔진으로,

NVIDIA GPU 가 탑재된 데스크탑 및 전용 엣지 디바이스에서 추론 성능을 극대화하는 기술입니다.

 

CUDA 를 지원하는 NVIDIA GPU 가 필수이며, 연산 방식을 최적화하여 추론 속도를 극도로 높인 것이 특징입니다.

 

- 추가로, 일전에 왜 비전 모델이 자연어 처리 모델보다 가벼운지를 질문받은 적이 있습니다.

당시에는 대답을 제대로 하지 못하여 그냥 구조상 그렇다고만 대답을 했었는데,

후에 생각해보니, 시각 정보를 처리하는 것은 인간과 같은 고차원적인 동물이 아니라 굉장히 작은 벌레수준 까지만 해도 처리가 가능한 영역인 반면, 언어나 사고와 같은 추상적이고 고차원적인 정보의 처리는 인간과 비슷한 수준의 두뇌 능력을 지녀야 할 수 있는 분야기 때문이라고 정리할 수 있었습니다. (요약하자면 시각 정보인 위치적 특징을 추출해 벡터로 만드는 작업)

이로인해 저는 현 시점에서는 제한된 컴퓨팅 환경에서도 다룰 수 있는 비전 모델이 더 매력적으로 생각하고 있으며,

모바일과 같은 작은 디바이스에서도 돌아가는 모델을 만드는 것을 목표로 개인 연구를 진행하고 있다고 생각합니다.