[Pytorch] pytorch로 학습해서 저장한 pth 파일을 cpu에 최적화 하는 방법

pytorch로 학습해서 저장한 pth 파일을 cpu에 최적화 하는 방법

PyTorch 모델을 CPU에 최적화하여 배포하기 위해, 모델을 TorchScript로 변환하고, 이를 통해 성능을 향상시킬 수 있습니다. TorchScript는 PyTorch 모델을 최적화하고, 이식 가능한 형식으로 변환하는 방법을 제공합니다. 이 과정을 통해, 모델을 다양한 플랫폼과 환경에서 실행할 수 있게 되며, 특히 CPU 환경에서의 성능을 개선할 수 있습니다.

---

1. TorchScript로 모델 변환

PyTorch 모델을 TorchScript로 변환하는 주요 방법에는 두 가지가 있습니다: 트레이싱(Tracing) 과 스크립팅(Scripting).

• 트레이싱(Tracing) 은 모델을 실행하면서 연산을 기록하여 TorchScript 모델을 생성합니다. 이 방법은 동적인 제어 흐름(dynamic control flow)을 포함하지 않는 모델에 적합합니다.
• 스크립팅(Scripting) 은 모델의 코드를 직접 분석하여 동적인 제어 흐름을 포함하는 모델을 TorchScript로 변환할 수 있게 합니다.

 

2. 모델 최적화

TorchScript 모델을 생성한 후에는 추가적인 최적화 단계를 수행할 수 있습니다. 예를 들어, 양자화(Quantization), 프루닝(Pruning), 모델 병합(Model Fusion) 등의 기술을 사용하여 모델의 크기를 줄이고, 실행 속도를 향상시킬 수 있습니다.

 

 

3. 예시 코드

모델을 TorchScript로 변환하기

import torch
import torchvision

# 예제 모델 로드
model = torchvision.models.resnet18(pretrained=True)
model.eval()

# 모델을 TorchScript로 트레이싱
example_input = torch.rand(1, 3, 224, 224)  # 모델 입력 크기에 맞게 조정
traced_script_module = torch.jit.trace(model, example_input)

# TorchScript 모델 저장
traced_script_module.save("traced_resnet18_model.pt")

 

양자화를 통한 모델 최적화

model_fp32 = torchvision.models.resnet18(pretrained=True)
model_fp32.eval()
model_fp32.qconfig = torch.quantization.get_default_qconfig('fbgemm')
model_fp32_fused = torch.quantization.fuse_modules(model_fp32, [['conv1', 'bn1', 'relu']])
model_int8 = torch.quantization.prepare(model_fp32_fused, inplace=False)
model_int8 = torch.quantization.convert(model_int8, inplace=False)

# 양자화된 모델 저장
model_int8.save("quantized_resnet18_model.pt")

 

이러한 과정을 거치면, PyTorch 모델을 CPU 환경에서 더 효율적으로 실행할 수 있게 최적화할 수 있습니다. 또한, 이 방법은 PyTorch 모델을 다양한 환경에서 배포하는 데 도움이 됩니다.