Rust + WASM 실전 프로젝트 시리즈 3

Rust + WASM 실전 프로젝트 시리즈 3: 실제 딥러닝 모델을 WASM에 포팅하여 보다 정밀한 인식을 구현

지난 시리즈에서는 Rust로 간단한 분류기나 룰 기반의 인식 시스템을 브라우저에 구현해보았습니다. 이번에는 그보다 한 단계 더 나아가, 실제 학습된 딥러닝 모델을 WASM으로 포팅하여 브라우저 환경에서도 보다 정밀한 인식 기능을 구현해보겠습니다.

WASM이 브라우저의 자바스크립트보다 훨씬 빠르게 수치 연산을 처리할 수 있다는 점은 이전에도 확인했지만, 이번 주제에서는 그 성능을 머신러닝 인퍼런스에 적극 활용해볼 계획입니다. 예제에서는 숫자 인식 또는 간단한 이미지 분류 모델을 ONNX 포맷으로 가져와 Rust에서 처리하고, 결과를 실시간으로 렌더링하는 방식으로 구성할 예정입니다.

목표

Python에서 학습한 딥러닝 모델을 ONNX 형식으로 저장
Rust + WASM 환경에서 모델을 불러와 실행
사용자의 입력 이미지에 대해 고정밀 예측 수행
브라우저 상에서 빠르고 실시간에 가까운 반응 제공

모델 준비: Python에서 ONNX 내보내기

먼저, Python에서 사전 학습된 모델을 ONNX 포맷으로 변환합니다. 예시로는 MNIST 숫자 분류기를 사용합니다.

import torch
import torchvision.models as models
import torch.onnx

model = models.mnist(pretrained=True)
dummy_input = torch.randn(1, 1, 28, 28)
torch.onnx.export(model, dummy_input, "mnist.onnx", opset_version=11)

이렇게 생성된 mnist.onnx 파일은 Rust 애플리케이션에서 사용할 수 있도록 tract 라이브러리를 통해 불러옵니다.

Rust에서 ONNX 모델 로딩

Rust에는 ONNX 런타임을 지원하는 다양한 라이브러리가 있지만, WASM 타겟과 호환이 잘 되는 것은 tract입니다. 특히 tract는 모델의 추론만 수행하는 lightweight 구조로, 브라우저 환경에서도 효율적으로 작동합니다.

[dependencies]
tract-onnx = "0.20"
wasm-bindgen = "0.2"

그리고 WASM에서 사용할 엔트리 포인트는 다음과 같이 정의할 수 있습니다.

use tract_onnx::prelude::*;
use wasm_bindgen::prelude::*;

#[wasm_bindgen]
pub fn predict(input: &[f32]) -> String {
    let model = load_model().unwrap();
    let result = run_inference(&model, input);
    result
}

load_model()은 static으로 빌드된 ONNX 모델을 포함하고 있으며, run_inference()는 입력 데이터를 텐서로 변환한 후 추론을 수행하고 결과를 반환합니다.

모델을 static으로 포함하기

브라우저에 배포하기 위해선 모델 파일을 WASM 바이너리에 포함해야 합니다. include_bytes! 매크로를 활용해 모델을 메모리로 직접 불러옵니다.

fn load_model() -> TractResult<SimplePlan<TypedFact, Box<dyn TypedOp>>> {
    let model_data = include_bytes!("../assets/mnist.onnx");
    let model = tract_onnx::onnx()
        .model_for_read(&mut &model_data[..])?
        .into_optimized()?
        .into_runnable()?;
    Ok(model)
}

브라우저와 연결: JS → WASM 호출

캔버스에서 이미지를 가져와 JS에서 grayscale → 1차원 벡터로 전처리한 후, Rust의 predict() 함수에 전달합니다. WASM에서 결과 값을 문자열로 반환하면 브라우저에 바로 출력할 수 있습니다.

import init, { predict } from './pkg/model_infer.js';

async function runPrediction(inputPixels) {
    await init();
    const result = predict(inputPixels);
    document.getElementById('result').textContent = `예측 결과: ${result}`;
}

이 구조는 성능을 JS와 WASM 간에 적절히 분리하면서도, 사용자는 웹 브라우저 상에서 마치 데스크탑 수준의 추론 응답 속도를 경험할 수 있게 만듭니다.

성능과 제약

WASM에서의 추론은 CPU 기반이며, GPU 가속은 WebGPU를 통해 별도 구현해야 합니다. 따라서 복잡한 모델(예: ResNet, BERT 등)은 속도 이슈가 있을 수 있지만, 간단한 분류기나 CNN 기반 모델은 충분히 실용적입니다.

초기 로딩 시 모델의 크기가 중요한 이슈입니다. ONNX 모델이 수 MB 이상일 경우, 페이지 로딩 시간이 증가하므로 브라우저에서 사용할 모델은 꼭 최적화된 형태로 준비하는 것이 좋습니다. quantization이나 구조 단순화를 통해 모델 크기를 줄이는 작업도 고려해야 합니다.

마무리

이번 글에서는 실제 딥러닝 모델을 Rust + WASM 환경에서 로딩하고 추론까지 실행해보는 전 과정을 다뤄봤습니다. 기존 자바스크립트만으로는 어려웠던 정밀한 인식 기능을, 이제는 브라우저 내에서도 충분히 구현할 수 있게 되었습니다.

앞으로는 이미지 인식 외에도 음성, 문장, 행동 등 다양한 도메인에서 이런 방식의 WASM 기반 인퍼런스가 더욱 중요해질 것으로 보입니다. 다음 글에서는 모델 경량화와 WebGPU 기반 추론도 함께 살펴볼 예정이니 관심 있는 분들은 계속 지켜봐 주세요.

rachel414 님의 블로그

Rust + WASM 실전 프로젝트 시리즈 3