
Rust로 MNIST를 해보니 Trait이 Batcher의 설계도가 되었던 이야기
요약
Rust의 머신러닝 프레임워크인 burn을 사용하여 MNIST 데이터셋을 처리하는 과정을 다룹니다. Trait을 통해 데이터 파이프라인의 설계도를 구현하는 방법과 Rust의 강력한 타입 시스템을 활용한 개발 경험을 공유합니다.
핵심 포인트
- Rust의 Trait을 활용하여 Batcher 인터페이스를 설계도처럼 구현 가능
- burn 프레임워크를 통한 효율적인 머신러닝 데이터 파이프라인 구축
- 컴파일러와 AI를 활용하여 Rust의 엄격한 타입 시스템에 적응하는 방법
서론
Rust Book을 다 읽고 나서, "다음에는 무언가 만들어보자"라는 생각에 머신러닝 프레임워크인 burn을 만져보기로 했다.
[dependencies]
burn = { version = "0.16", features = ["ndarray"] }
Cargo.toml에 몇 줄 적고 cargo build를 실행하니 간단히 작동한다.
Rust에 이렇게나 많은 crate가 있다니 하는 것이 첫 번째 놀라움이었다.
PyTorch나 TensorFlow 같은 프레임워크가 Rust에도 존재하며, 심지어 제대로 유지보수되고 있다.
"바퀴를 다시 발명할 필요가 없다"라는 감각은 Rust Book을 읽는 것만으로는 얻을 수 없었던 것이었다.
Trait은 설계도였다
Rust Book 안에서 Trait은 "인터페이스와 같은 것"이라고 설명된다.
읽었을 때는 대충 이해한 것 같았지만, 솔직히 와닿지는 않았다.
burn의 데이터 파이프라인을 구현할 때, Batcher라는 Trait을 구현해야 했다.
impl<B: Backend> Batcher<B, MnistItem, MnistBatch<B>> for MnistBatcher {
fn batch(&self, items: Vec<MnistItem>, device: &B::Device) -> MnistBatch<B> {
// 이 부분을 직접 채운다
...
이 시그니처(signature)를 보았을 때, "아, 이것은 설계도다"라고 생각했다.
- 무엇을 받는가 →
Vec<MnistItem> - 무엇을 반환하는가 →
MnistBatch<B> - 어떤 환경에서 동작하는가 →
B: Backend(CPU·GPU 등을 추상화하고 있음)
Trait이 "무엇을 만들어야 하는지"를 알려주고 있다.
남은 것은 내용을 채우는 것뿐이다. burn이 준비한 틀에 자신의 처리를 끼워 넣는 감각이었다.
batch() 내부를 들여다보기
구현한 batch()의 내용은 다음과 같은 처리가 되어 있다.
fn batch(&self, items: Vec<MnistItem>, device: &B::Device) -> MnistBatch<B> {
let images = items
.iter()
...
흐름을 정리하면 다음과 같다.
생데이터 [[f32; 28]; 28]
→ TensorData로 변환 (부동 소수점으로 타입 변환)
→ Tensor<B, 2> [28, 28]를 생성
...
정규화의 0.1307 (평균)과 0.3081 (표준 편차)은 PyTorch의 MNIST 공식 샘플에서 사용되는 값을 그대로 사용하고 있다.
코멘트에 URL까지 적혀 있었다.
crate를 사용하면 이러한 "관례"도 함께 따라온다.
스스로 처음부터 조사하지 않아도 선인들의 지혜가 코드에 심어져 있다.
컴파일러와 AI를 두 명의 스승으로 삼기
Python이나 Jupyter notebook에 익숙하다면, Rust의 개발 경험은 처음에 다가가기 어렵다.
Python이라면 이렇게 즉시 확인할 수 있다.
# notebook에서 실행하여 확인
batch.images.shape # → (3, 28, 28)
Rust에서는 그렇게 되지 않는다. 컴파일하고 실행해야 비로소 값을 볼 수 있다.
하지만 학습 루프는 이렇게 바뀌었다.
코드를 작성한다
→ cargo build로 컴파일
→ 에러가 발생한다 (Rust의 에러 메시지는 친절하고 상세하다)
...
컴파일러는 "무엇이 틀렸는지"를 정확하게 알려준다.
AI는 "왜 그렇게 되는지", "Rust답게 어떻게 생각해야 하는지"를 설명해 준다.
이 두 사람을 스승으로 삼아 써 내려가면, notebook이 없어도 의외로 앞으로 나아갈 수 있다.
한 가지 보충하자면, Rust의 타입 시스템(type system)은 상당히 강력해서, Python의 notebook에서 실행해 보고 나서야 비로소 깨닫게 되는 많은 에러를 컴파일 시점에 걸러내 준다.
"컴파일이 통과하면 절반은 성공한 것이나 다름없다"라는 감각이 있다.
Rust에서 테스트를 작성하는 의미
그럼에도 불구하고, 컴파일이 통과한 후에 확인이 필요한 부분이 있다.
Tensor의 shape가 그 전형적인 예다.
pub images: Tensor<B, 3>,
Tensor<B, 3>
Tensor<B, 3>
라는 타입은 "3차원 텐서이다"라는 점은 보장해 준다.
하지만 각 차원의 크기([N, 28, 28]에서 N이나 28 부분)는 실행 시점에만 알 수 있다.
그래서 다음과 같은 테스트를 작성한다.
#[test]
fn test_batch_shapes() {
let batcher = MnistBatcher::default();
...
}
Python의 notebook에서 shift+Enter를 하던 부분이 Rust에서는 테스트가 된다.
게다가 이 테스트는 cargo test로 언제든 재실행할 수 있다.
notebook의 셀이 사라져도, 테스트는 남는다.
한 번 작성해 두면, 코드를 변경할 때마다 "이전과 동일하게 동작하는지"를 확인해 준다.
처음에는 테스트를 작성하는 것이 번거롭게 느껴졌지만, 지금은 "실험의 기록"으로서 작성하게 되었다.
요약
data.rs를 작성하며 배운 내용을 정리하면 다음과 같다.
- 풍부한 crate → 바퀴를 재발명할 필요가 줄어든다. 선배들의 지혜를 그대로 빌려올 수 있다.
- Trait이 설계도 → "무엇을 만들어야 하는지"가 시그니처 (Signature)에 적혀 있다. 나머지는 내용을 채우기만 하면 된다.
- 컴파일러와 AI → 이 두 명을 스승으로 삼으면 notebook이 없어도 앞으로 나아갈 수 있다.
- 테스트가 실험 기록 → notebook의
shift+Enter를 대신한다. 게다가 계속 남아있다.
Rust Book에서 배운 내용이 실제 코드 속에서 연결되는 순간이 있다.
impl Trait for Struct라는 구문이 단순한 문법이 아니라, "설계도를 받아 구현하는" 행위로 보이기 시작한 것이 바로 이 파일을 작성했을 때였다.
Rust Book을 읽다 보면 Trait이나 타입 시스템 (Type System)이 다소 추상적으로 보일 수 있다.
하지만 실제 라이브러리를 사용해 보면, 그것들이 "안전하게 확장하기 위한 설계도"라는 것을 알 수 있다.
Burn의 data.rs를 작성함으로써, 아주 조금이지만 Rust의 문법이 실전 도구로서 보이기 시작했다.
저자의 다른 실험
LiDAR 점군 (Point Cloud) 분석, 블록체인 × AI 등 다른 프로젝트들도 모아두었습니다.
Discussion

AI 자동 생성 콘텐츠
본 콘텐츠는 Zenn AI의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기