TorchScript로 모델 내보내기

PyTorch Lightning으로 만드는 확장 가능한 AI 모델

Sergiy Tkachuk

Director, GenAI Productivity

TorchScript란?

Python과 독립적
프로덕션에서 효율적
예:
- 모바일 기기 배포

휴대폰

TorchScript란?

Python과 독립적
프로덕션에서 효율적
예:
- 모바일 기기 배포
- 프로덕션 시스템에서 고성능 추론

모바일폰과 기어가 있는 노트북: 프로덕션 환경을 나타냄

모델을 TorchScript로 변환

변환 방법 2가지:

torch.jit.trace: 예시 입력으로 실행 경로 추적
torch.jit.script: Python 코드를 분석해 컴파일

사용 기준:

단순 모델엔 trace 사용
제어 흐름(예: 루프)이 있으면 script 사용

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def forward(self, x):
        return x * 2

model = SimpleModel()
scripted_model = torch.jit.script(model)

TorchScript 모델 저장과 로드

모델 저장:
- torch.jit.save: 스크립트된 모델을 파일로 저장
모델 로드:
- torch.jit.load: 추론을 위해 모델 로드

# Save the model
torch.jit.save(scripted_mod,"model.pt")

# Load the model
loaded_model=torch.jit.load("model.pt")

TorchScript로 추론 수행

단계:
- TorchScript 모델 로드
- 입력을 전달해 예측 수행
- 출력은 PyTorch 예측과 동일

입력 예:

입력 텐서: [1.0, 2.0, 3.0]

출력 예:

출력 텐서: [2.0, 4.0, 6.0]

# Perform inference
input_arr = [1.0, 2.0, 3.0]
input_tensor = torch.tensor(input_arr)

output = loaded_model(input_tensor)
print(output)

TorchScript 한눈에 보기

torch.jit.trace: 정적 모델에 적합
torch.jit.script: 동적 제어 흐름 처리
torch.jit.save: 스크립트된 모델 저장
torch.jit.load: 추론용으로 다시 로드

Ayo berlatih!

PyTorch Lightning으로 만드는 확장 가능한 AI 모델