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

모델 경량화와 WebGPU 기반 추론

이전 글에서는 ONNX 모델을 Rust에서 WASM으로 포팅하고, 브라우저에서 추론하는 과정을 살펴보았습니다. 성능과 정확도는 꽤 만족스러웠지만, 실제 배포를 고려하면 몇 가지 한계점이 존재합니다. 특히 모델 크기와 실행 속도, 그리고 연산 효율성이 주요한 병목으로 작용할 수 있습니다.

이번 글에서는 두 가지에 집중합니다. 첫째는 모델 경량화, 둘째는 WebGPU 기반 추론 구조</strong입니다. 두 가지 모두 브라우저 환경에서의 AI 인퍼런스를 보다 실용적으로 만들기 위한 핵심 기술입니다.

모델 경량화: 용량과 연산 최소화

ONNX나 TensorFlow 모델을 WASM 환경에 올릴 경우, 가장 먼저 마주하게 되는 문제가 모델 파일의 크기입니다. 일반적인 CNN 모델조차 수 MB~수십 MB를 차지하며, 이는 초기 로딩 시간에 직접적인 영향을 미칩니다. 따라서 가능한 한 모델을 가볍게 만들어야 합니다.

1. 양자화 (Quantization)

양자화는 32-bit float 대신 8-bit 정수 등을 사용하는 방식으로, 정확도 손실을 최소화하면서도 연산 속도를 높이고 모델 크기를 줄입니다. PyTorch 기준 예시는 다음과 같습니다.

import torch
from torch.quantization import quantize_dynamic

model_fp32 = torch.load("mnist_model.pth")
model_int8 = quantize_dynamic(model_fp32, {torch.nn.Linear}, dtype=torch.qint8)
torch.save(model_int8, "mnist_quantized.pth")

이렇게 생성한 모델은 ONNX로 내보낼 수 있으며, 크기가 약 4~5배 이상 줄어드는 효과가 있습니다. 실제 추론 시간도 약간 줄어들 수 있습니다.

2. 구조 단순화 및 슬림 아키텍처

복잡한 구조의 모델은 정확도는 높을 수 있지만 모바일이나 브라우저 환경엔 무리입니다. 따라서 모델 구조를 MobileNet, EfficientNet-lite 등으로 바꾸는 것이 좋습니다.

또한 채널 수나 필터 수를 줄이는 것만으로도 상당한 경량화가 가능하며, 이 과정에서 정확도 손실을 테스트하며 균형을 맞추는 것이 중요합니다.

WebGPU 기반 추론

기존 WASM 기반 추론은 CPU 연산에만 의존하기 때문에, 복잡한 연산에서는 속도나 효율 면에서 한계를 보입니다. 특히 고해상도 이미지 처리나 CNN 구조에서는 병렬처리가 필수입니다.

WebGPU는 이런 병목을 해결할 수 있는 강력한 대안입니다. 브라우저 내에서 GPU 자원을 직접 활용할 수 있으며, WASM과도 자연스럽게 연동이 가능합니다.

1. WebGPU란?

WebGPU는 최신 웹 브라우저에서 GPU를 사용할 수 있게 해주는 API로, WebGL의 후속 기술입니다. Vulkan, Metal, DirectX12의 공통 기반으로 설계되었고, JavaScript뿐 아니라 WASM에서도 호출이 가능합니다.

즉, Rust에서 WebGPU 라이브러리를 사용해 GPU 연산을 구성하고, WASM 빌드로 브라우저에서 이를 실행하는 구조가 가능합니다.

2. Rust에서 WebGPU 사용하기

Rust에서는 wgpu라는 라이브러리를 통해 WebGPU API에 접근할 수 있습니다. 예를 들어, 2D 텐서 곱셈을 GPU에서 실행하는 코드는 다음과 같이 구성할 수 있습니다.

[dependencies]
wgpu = "0.19"
pollster = "0.3"
wasm-bindgen = "0.2"

use wgpu::util::DeviceExt;

async fn run_gpu_inference(input: &[f32], weights: &[f32]) -> Vec<f32> {
    let instance = wgpu::Instance::default();
    let adapter = instance.request_adapter(&Default::default()).await.unwrap();
    let (device, queue) = adapter.request_device(&Default::default(), None).await.unwrap();

    // 버퍼 설정 및 compute shader 실행 생략
    // 실제 커널은 WGSL 또는 SPIR-V로 작성
}

WGSL(웹GPU 셰이더 언어)을 이용해 CNN의 conv 연산이나 dense layer 등을 정의할 수 있으며, Rust에서 이 커널을 호출하여 고속 병렬 추론을 실행할 수 있습니다.

3. WASM + WebGPU 연동

현재 WebGPU는 일부 브라우저(Chrome, Edge, Firefox Nightly 등)에서만 기본 활성화되어 있지만, 실험적인 기능으로 충분히 사용할 수 있습니다.

WebGPU API는 JS에서 먼저 초기화한 후, SharedBuffer 또는 WASM 메모리를 통해 Rust와 연결하는 방식으로 구성됩니다. 구조 예시는 다음과 같습니다.

// JavaScript: WebGPU 초기화
const adapter = await navigator.gpu.requestAdapter();
const device = await adapter.requestDevice();
// 이후 buffer와 compute pipeline 구성

그리고 Rust에서 해당 버퍼에 접근하거나, 별도로 GPU 자원을 요청하여 연산에 활용합니다. 이 과정은 아직 표준화가 완전히 끝난 상태는 아니므로 구현 시 브라우저 호환성을 반드시 확인해야 합니다.

결론

Rust + WASM으로 구현된 인퍼런스 시스템은 경량화된 모델과 WebGPU의 결합을 통해 실제 애플리케이션 수준의 반응성과 정확도를 제공할 수 있습니다.

모델 경량화는 추론 속도를 높이고 초기 로딩을 단축하며, WebGPU는 복잡한 연산을 빠르게 처리하게 해줍니다. 이 두 가지가 조화를 이루면 브라우저에서 실행되는 AI 서비스도 본격적인 수준에 도달하게 됩니다.

다음 글에서는 WebAssembly System Interface(WASI)를 활용해 브라우저가 아닌 서버리스 환경에서 WASM AI를 실행하는 구조를 소개할 예정입니다. 브라우저에서 학습은 어렵지만, 추론은 점점 더 다양한 영역으로 확장되고 있습니다.

rachel414 님의 블로그