4장: PyTorch TorchScript

PyTorch 모델의 프로덕션 배포

弘益人間 · 널리 인간을 이롭게 하라

4.1 TorchScript 소개

TorchScript는 PyTorch 코드에서 직렬화 가능하고 최적화 가능한 모델을 생성하는 PyTorch의 방법입니다. 연구(Python)와 프로덕션(C++) 간의 격차를 해소하여 Python 종속성 없이 환경에 배포할 수 있습니다. TorchScript는 JIT(Just-In-Time) 컴파일을 사용하여 추론을 위한 모델을 최적화합니다.

TorchScript의 핵심 철학은 Python의 동적 즉시 실행에서 프로덕션 환경의 정적 최적화 실행으로 전환하는 경로를 제공하는 것입니다. 이러한 전환은 성능, 안정성, 최소 종속성이 중요한 대규모 모델 배포에 매우 중요합니다.

4.2 Tracing vs Scripting

TorchScript는 PyTorch 모델을 변환하는 두 가지 방법을 제공합니다: 추적(tracing)과 스크립팅(scripting). 각 방법을 언제 사용할지 이해하는 것이 성공적인 모델 배포에 중요합니다.

Tracing (추적)

추적은 예제 입력으로 순방향 패스 중에 실행되는 연산을 기록합니다. 정적 계산 그래프를 가진 모델에 이상적입니다.

import torch

# 모델 정의
class SimpleModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = torch.nn.Linear(10, 5)
        self.relu = torch.nn.ReLU()

    def forward(self, x):
        return self.relu(self.fc(x))

model = SimpleModel()
example_input = torch.randn(1, 10)

# 모델 추적
traced_model = torch.jit.trace(model, example_input)

# 저장
traced_model.save('traced_model.pt')

# 로드 및 사용
loaded = torch.jit.load('traced_model.pt')
output = loaded(torch.randn(1, 10))

# 추적된 그래프 검사
print(traced_model.graph)
print(traced_model.code)

Scripting (스크립팅)

스크립팅은 Python 소스 코드를 직접 분석하여 제어 흐름과 동적 동작을 캡처합니다. 모델에 if 문, 루프 또는 기타 제어 구조가 포함된 경우 사용합니다.

class ConditionalModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = torch.nn.Linear(10, 5)

    def forward(self, x):
        # 보존해야 하는 제어 흐름
        if x.sum() > 0:
            return self.fc(x)
        else:
            return x[:, :5]

model = ConditionalModel()

# 모델 스크립팅 (제어 흐름 캡처)
scripted_model = torch.jit.script(model)

# 저장
scripted_model.save('scripted_model.pt')

# 제어 흐름이 보존되었는지 확인
print(scripted_model.code)

4.3 Tracing vs Scripting 선택 기준

모델의 특성에 따라 적절한 방법을 선택하세요:

4.4 고급 TorchScript 기능

타입 어노테이션

타입 어노테이션은 스크립팅에 필수적이며, TorchScript가 코드의 의도를 이해하고 조기에 오류를 발견하는 데 도움이 됩니다.

import torch
from typing import List, Tuple, Dict, Optional

class AnnotatedModel(torch.nn.Module):
    def forward(
        self,
        x: torch.Tensor,
        hidden: Tuple[torch.Tensor, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        h, c = hidden
        output = self.process(x, h)
        return {"output": output, "hidden": c}

    def process(self, x: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
        return x + h

    def batch_process(
        self,
        inputs: List[torch.Tensor],
        mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """선택적 마스킹으로 배치 처리"""
        result = torch.stack(inputs)
        if mask is not None:
            result = result * mask
        return result

사용자 정의 연산자

import torch

@torch.jit.script
def custom_activation(x: torch.Tensor, threshold: float) -> torch.Tensor:
    """사용자 정의 게이트 활성화 함수"""
    return torch.where(x > threshold, x, torch.zeros_like(x))

@torch.jit.script
def swish(x: torch.Tensor, beta: float = 1.0) -> torch.Tensor:
    """Swish 활성화: x * sigmoid(beta * x)"""
    return x * torch.sigmoid(beta * x)

class ModelWithCustomOp(torch.nn.Module):
    def forward(self, x):
        x = custom_activation(x, 0.5)
        x = swish(x, beta=1.5)
        return x

model = ModelWithCustomOp()
scripted = torch.jit.script(model)

4.5 TorchScript 모델 로드 및 사용

Python에서

import torch
import numpy as np

# 모델 로드
model = torch.jit.load('model.pt')
model.eval()

# 추론 실행
with torch.no_grad():
    output = model(input_tensor)

# 그래프 검사
print("그래프 구조:")
print(model.graph)

print("\n생성된 코드:")
print(model.code)

# 모델 파라미터 확인
for name, param in model.named_parameters():
    print(f"{name}: {param.shape}")

# 추론 벤치마크
import time
start = time.time()
for _ in range(100):
    with torch.no_grad():
        _ = model(input_tensor)
elapsed = (time.time() - start) / 100
print(f"평균 추론 시간: {elapsed*1000:.2f}ms")

C++에서

#include <torch/script.h>
#include <iostream>
#include <memory>
#include <chrono>

int main() {
    // 모델 로드
    torch::jit::script::Module module;
    try {
        module = torch::jit::load("model.pt");
        module.eval();
    } catch (const c10::Error& e) {
        std::cerr << "모델 로드 오류: " << e.what() << "\n";
        return -1;
    }

    // 입력 생성
    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(torch::randn({1, 3, 224, 224}));

    // 타이밍과 함께 모델 실행
    auto start = std::chrono::high_resolution_clock::now();
    at::Tensor output = module.forward(inputs).toTensor();
    auto end = std::chrono::high_resolution_clock::now();

    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
    std::cout << "추론 시간: " << duration.count() << "ms\n";
    std::cout << "출력 형상: " << output.sizes() << '\n';
    std::cout << "출력 (처음 5개): " << output.slice(1, 0, 5) << '\n';

    return 0;
}

// 컴파일:
// g++ -std=c++17 inference.cpp -o inference \
//     -I${TORCH_PATH}/include \
//     -L${TORCH_PATH}/lib \
//     -ltorch -ltorch_cpu -lc10

4.6 최적화 기법

그래프 최적화

import torch

model = torch.jit.load('model.pt')

# 추론을 위한 최적화
optimized_model = torch.jit.optimize_for_inference(model)

# 모델 동결 (상수 인라인화, 훈련 전용 연산 제거)
frozen_model = torch.jit.freeze(optimized_model)

# 추가 최적화
# 1. 드롭아웃 레이어 제거
frozen_model.eval()

# 2. 연산 융합
# Conv + BatchNorm + ReLU가 자동으로 융합됨

# 최적화된 버전 저장
frozen_model.save('optimized_model.pt')

# 모델 크기 비교
import os
original_size = os.path.getsize('model.pt') / (1024 * 1024)
optimized_size = os.path.getsize('optimized_model.pt') / (1024 * 1024)
print(f"원본: {original_size:.2f}MB")
print(f"최적화: {optimized_size:.2f}MB")
print(f"감소율: {(1 - optimized_size/original_size)*100:.1f}%")

연산자 융합

PyTorch는 메모리 대역폭을 줄이고 성능을 향상시키기 위해 일반적인 연산자 패턴을 자동으로 융합합니다.

# PyTorch는 자동으로 연산 융합
# Conv + BatchNorm + ReLU → 융합 연산

model = torch.nn.Sequential(
    torch.nn.Conv2d(3, 64, 3),
    torch.nn.BatchNorm2d(64),
    torch.nn.ReLU()
)

traced = torch.jit.trace(model, torch.randn(1, 3, 224, 224))
# 추적 중 자동으로 융합 발생

# 일반적인 융합 패턴:
# 1. Conv + BN + ReLU
# 2. Linear + ReLU
# 3. Add + ReLU
# 4. Mul + Add (fused multiply-add)

# 융합된 그래프 보기
print(traced.graph)

4.7 모바일 배포

PyTorch Mobile

import torch
from torch.utils.mobile_optimizer import optimize_for_mobile

model = torch.jit.load('model.pt')
model.eval()

# 모바일 최적화
mobile_model = optimize_for_mobile(model)

# 라이트 인터프리터용 저장 (더 작은 런타임)
mobile_model._save_for_lite_interpreter('mobile_model.ptl')

# 모바일 모델 검증
loaded_mobile = torch.jit.load('mobile_model.ptl')
test_input = torch.randn(1, 3, 224, 224)
output = loaded_mobile(test_input)
print(f"모바일 모델 출력 형상: {output.shape}")

# Android에서 (Java):
"""
import org.pytorch.LiteModuleLoader;
import org.pytorch.Module;
import org.pytorch.Tensor;

Module module = LiteModuleLoader.load("mobile_model.ptl");
Tensor input = Tensor.fromBlob(inputArray, shape);
Tensor output = module.forward(IValue.from(input)).toTensor();
float[] scores = output.getDataAsFloatArray();
"""

# iOS에서 (Swift):
"""
import LibTorch

guard let module = try? LiteModuleLoader.load(modelPath: "mobile_model.ptl") else {
    fatalError("모델 로드 실패")
}
let input = try! Tensor(shape: [1, 3, 224, 224])
let output = try! module.forward([input])
"""

4.8 양자화

양자화는 가중치와 활성화를 부동소수점에서 정수로 변환하여 모델 크기를 줄이고 추론 속도를 향상시킵니다.

import torch

# 후처리 정적 양자화
model.eval()

# 1단계: 모듈 융합
model_fused = torch.ao.quantization.fuse_modules(
    model,
    [['conv', 'bn', 'relu']]
)

# 2단계: 양자화 구성 지정
model_fused.qconfig = torch.ao.quantization.get_default_qconfig('x86')

# 모바일 배포의 경우:
# model_fused.qconfig = torch.ao.quantization.get_default_qconfig('qnnpack')

# 3단계: 양자화 준비
model_prepared = torch.ao.quantization.prepare(model_fused)

# 4단계: 샘플 데이터로 보정
with torch.no_grad():
    for data, _ in calibration_loader:
        model_prepared(data)

# 5단계: 양자화된 모델로 변환
quantized_model = torch.ao.quantization.convert(model_prepared)

# 양자화된 모델 저장
torch.jit.save(torch.jit.script(quantized_model), 'quantized.pt')

# 크기 및 성능 비교
original_size = sum(p.numel() * p.element_size() for p in model.parameters())
quantized_size = sum(p.numel() * p.element_size() for p in quantized_model.parameters())
print(f"크기 감소: {original_size / quantized_size:.1f}배")

# 벤치마크
import time
inputs = torch.randn(10, 3, 224, 224)

start = time.time()
with torch.no_grad():
    _ = model(inputs)
fp32_time = time.time() - start

start = time.time()
with torch.no_grad():
    _ = quantized_model(inputs)
int8_time = time.time() - start

print(f"속도 향상: {fp32_time / int8_time:.1f}배")

4.9 모범 사례

동적 형상 처리

class DynamicModel(torch.nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size = x.size(0)
        seq_len = x.size(1)
        # 동적 차원을 적절하게 사용
        return self.process(x, batch_size, seq_len)

    def process(self, x: torch.Tensor, bs: int, seq: int) -> torch.Tensor:
        # 하드코딩된 형상 피하기
        output = torch.zeros(bs, seq, 512, device=x.device)
        return output + x

# 다양한 입력 크기로 추적하여 검증
traced = torch.jit.trace(model, torch.randn(1, 10, 512))

# 다양한 배치 크기로 테스트
assert traced(torch.randn(4, 20, 512)).shape[0] == 4
assert traced(torch.randn(8, 15, 512)).shape[0] == 8

# 진정한 동적 동작을 위해 torch.jit.script 사용
scripted = torch.jit.script(model)

오류 처리

try:
    model = torch.jit.script(MyModel())
    print("스크립팅 성공!")
except RuntimeError as e:
    print(f"스크립팅 실패: {e}")
    print("대신 추적 시도 중...")
    try:
        example_input = torch.randn(1, 3, 224, 224)
        model = torch.jit.trace(MyModel(), example_input)
        print("추적 성공!")
    except Exception as e:
        print(f"스크립팅과 추적 모두 실패: {e}")
        # 모델 코드 수정 또는 즉시 모드 사용 고려

# 변환된 모델 검증
def validate_model(original, converted, test_inputs):
    """변환된 모델이 원본과 일치하는지 검증"""
    original.eval()
    converted.eval()

    with torch.no_grad():
        for inp in test_inputs:
            orig_out = original(inp)
            conv_out = converted(inp)
            diff = torch.abs(orig_out - conv_out)
            max_diff = torch.max(diff).item()
            print(f"최대 차이: {max_diff:.2e}")
            assert max_diff < 1e-5, f"출력이 너무 다름: {max_diff}"

test_inputs = [torch.randn(1, 3, 224, 224) for _ in range(5)]
validate_model(original_model, traced_model, test_inputs)

4.10 TorchScript 모델 디버깅

TorchScript 모델 디버깅에는 표준 Python 디버깅과 다른 기술이 필요합니다.

생성된 코드 검사

import torch

# 스크립팅된 모델 로드 또는 생성
model = torch.jit.script(MyModel())

# 생성된 TorchScript 코드 보기
print(model.code)

# 계산 그래프 보기
print(model.graph)

# 상세한 그래프 표현 가져오기
for node in model.graph.nodes():
    print(f"노드: {node.kind()}")
    print(f"  입력: {[i.debugName() for i in node.inputs()]}")
    print(f"  출력: {[o.debugName() for o in node.outputs()]}")

# 특정 연산 찾기
def find_operations(graph, op_name):
    """특정 연산 유형의 모든 노드 찾기"""
    return [node for node in graph.nodes() if node.kind() == op_name]

conv_ops = find_operations(model.graph, 'aten::conv2d')
print(f"발견된 컨볼루션 연산: {len(conv_ops)}개")

성능 프로파일링

import torch
from torch.profiler import profile, ProfilerActivity

model = torch.jit.load('model.pt')
input_data = torch.randn(1, 3, 224, 224)

# PyTorch 프로파일러로 프로파일링
with profile(
    activities=[ProfilerActivity.CPU],
    record_shapes=True,
    profile_memory=True
) as prof:
    with torch.no_grad():
        model(input_data)

# 프로파일링 결과 출력
print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))

# Chrome 추적 내보내기 (시각화용)
prof.export_chrome_trace("trace.json")
# chrome://tracing에서 trace.json 열기

실전 배포 팁

import torch

# 1. 배치 크기별 성능 측정
def benchmark_batch_sizes(model, input_shape, batch_sizes=[1, 4, 8, 16]):
    """다양한 배치 크기로 모델 벤치마크"""
    results = {}
    for bs in batch_sizes:
        input_data = torch.randn(bs, *input_shape[1:])

        # 웜업
        for _ in range(10):
            with torch.no_grad():
                _ = model(input_data)

        # 측정
        import time
        start = time.time()
        for _ in range(100):
            with torch.no_grad():
                _ = model(input_data)
        elapsed = (time.time() - start) / 100

        results[bs] = {
            'latency_ms': elapsed * 1000,
            'throughput': bs / elapsed
        }

    return results

# 2. 메모리 사용량 프로파일링
def profile_memory(model, input_tensor):
    """모델 메모리 사용량 측정"""
    import torch.cuda as cuda

    if cuda.is_available():
        cuda.reset_peak_memory_stats()
        with torch.no_grad():
            _ = model(input_tensor.cuda())
        peak_memory = cuda.max_memory_allocated() / (1024 ** 2)  # MB
        print(f"피크 메모리 사용량: {peak_memory:.2f}MB")
    else:
        print("CUDA를 사용할 수 없습니다")

# 3. 다중 스레드 추론 테스트
def test_multithreaded_inference(model, num_threads=4):
    """다중 스레드 환경에서 모델 테스트"""
    import concurrent.futures

    def inference_task(idx):
        input_data = torch.randn(1, 3, 224, 224)
        with torch.no_grad():
            output = model(input_data)
        return output.shape

    with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor:
        futures = [executor.submit(inference_task, i) for i in range(10)]
        results = [f.result() for f in futures]

    print(f"✓ {num_threads}개 스레드로 다중 스레드 추론 성공")