DGX SparkでRustを使ったディープラーニング！

はじめに

こんにちは！トグルホールディングスでエンジニアをしている木下隼です。
toggle Advent Calendar 2025の22日目となります！

今回、弊社のエンジニア全員にDGX Sparkが配備されました！実際の活用例として技術創発チームでは、ディープラーニングを活用しているケースもあります。現在だと、インターン生がYOLOを使った画像認識モデルの学習をDGX Sparkで行なっていたりもします。

また、弊チームではRustを使うエンジニアが多く...自分自身がRustの勉強をしたいと思っていました。機械学習は自分の好きな分野ということもあり、機械学習を通じてRustを学んでいます。

さらにDGX SparkでBurnなどを使った記事が少なかったため、何かしらの知見となればと思い、挑戦してみました！

Rustの深層学習フレームワーク「Burn」とは

Burnは、Rustで書かれた深層学習フレームワークです。PyTorchライクな構成で、Tensor演算・自動微分・学習ループ（train）などが一通り扱えます。

PyTorch/NumPyは内部実装の多くがC/C++で、Pythonからそれらを呼び出す形ですが、BurnはコアがRustで実装されています（ただしバックエンドによってはネイティブAPIを使用する場合あり）。
※本記事では Burn 0.19.1 を使っています。

GPU環境が対応しているのか？

Burnはバックエンドを差し替えることで、CPU/GPUや実行環境を簡単に切り替えられます。

• NdArray：CPUバックエンド
• Cuda：NVIDIA GPU向けのCUDAバックエンド
• Wgpu：wgpu（WebGPU）経由で Vulkan/Metal/DX12 などに載せるバックエンド

Wgpuは幅広いGPU/OSで動かしやすいのが魅力で、CudaはNVIDIA環境で安定・高速になりやすい、というイメージです（体感）。

Rust✖️Burnで実装する利点

個人的に書いていていいなと思ったのは以下の点です。

• Rustの型システムのおかげで、型や一部の不整合を早めに潰しやすい（ただしOOMや環境依存、入力データの不整合など実行時エラーは起こり得ます）
• GPU・CPUバックエンドの切替を、型（ジェネリクス）で素直に表現できる
• モジュールやトレイトによる実装で拡張しやすい
• メソッドチェーンで処理を組み立てられて読みやすい
• 学習中の進捗表示（TUI）が見やすい（個人的に好きなポイント）

DGX Sparkで動かしてみる

今回はCNNモデルを構築し、建築図面の家具分類タスクを学習させてみます。

Cargo.toml

Burn周りをインストールします。

Cargo.toml

[package]
name = "burn_mnist"
version = "0.1.0"
edition = "2021"

[dependencies]
# Burn本体
burn = { version = "0.19.1", features = ["train", "ndarray", "vision", "wgpu", "cuda"] }

# シリアライズ（モデル保存）
serde = { version = "1.0", features = ["derive"] }

データ部分の実装

今回の使用データはこちらです： Furnishing Dataset (Kaggle)

複数の画像サイズが含まれていたので、事前に画像サイズを224x224に揃えました。
また、この実装では data/FDS/train / data/FDS/test 配下に「クラス名フォルダ」を並べる構成が前提になります。
（このデータセット自体この構成になっているのでそのまま使えます）

データのロード

データをロードするためのデータローダーを作成します。まずはBurn関連のモジュールをインポートします。

fds/data.rs

use burn::{
    data::{
        dataloader::batcher::Batcher,
        dataset::vision::{
            Annotation, ImageDatasetItem, ImageFolderDataset, PixelDepth
        },
    },
    prelude::*
};

ローダーの実装は、トレイトで実装しています。こうしておくことで、ImageFolderDatasetという型に対して、FdsLoaderトレイトを追加実装する形にでき、今回のデータを扱うときはこれをインポートして、fds_trainを呼び出せばデータが取得できる状態にできます。

fds/data.rs

pub trait FdsLoader {
    fn fds_train() -> Self;
    fn fds_test() -> Self;
}

impl FdsLoader for ImageFolderDataset {
    fn fds_train() -> Self {
        Self::new_classification("data/FDS/train").unwrap()
    }

    fn fds_test() -> Self {
        Self::new_classification("data/FDS/test").unwrap()
    }
}

また、ImageFolderDataset::new_classificationはとても便利な関数で、以下のような構造のデータディレクトリを指定すると、自動でラベル付きのデータを作成してくれます。

your_dataset/
├── train/
│   ├── class1/
│   │   ├── image1.jpg
│   │   ├── image2.jpg
│   │   └── ...
│   ├── class2/
│   │   ├── image1.jpg
│   │   └── ...
│   └── ...
└── test/
    ├── class1/
    │   └── ...
    └── class2/
        └── ...

バッチ処理用の構造体を実装

PyTorchではバッチ処理の内部実装などは基本的にPyTorchのモジュールを使えば簡単にできていました。Burnではこの部分はちゃんと自分で実装する必要があります。
まず、FdsBatchというバッチデータの構造体を定義し、データの構造を決めます。
そして属性なしのFdsBatcherを作成し、newメソッドでバッチ処理を行うオブジェクトを作成できる箱を作ります。

fds/data.rs

#[derive(Clone, Debug)]
pub struct FdsBatch<B: Backend> {
    pub images: Tensor<B, 4>,
    pub targets: Tensor<B, 1, Int>,
}

#[derive(Clone)]
pub struct FdsBatcher;

impl FdsBatcher {
    pub fn new() -> Self {
        Self
    }
}

そしてBatcherという名前でデータに対して前処理やバッチデータに変換する部分を実装していきます。
画像の方は、次元を入れ替えるのと、バッチ分の次元を追加して正規化しています。

fds/data.rs

impl<B: Backend> Batcher<B, ImageDatasetItem, FdsBatch<B>> for FdsBatcher {
    fn batch(&self, items: Vec<ImageDatasetItem>, device: &B::Device) -> FdsBatch<B> {
        fn image_as_vec_u8(item: ImageDatasetItem) -> Vec<u8> {
            item.image
                .into_iter()
                .map(|p: PixelDepth| -> u8 { p.try_into().unwrap() })
                .collect::<Vec<u8>>()
        }

        let targets = items
            .iter()
            .map(|item| {
                if let Annotation::Label(y) = item.annotation {
                    Tensor::<B, 1, Int>::from_data(
                        TensorData::from([(y as i64).elem::<B::IntElem>()]),
                        device,
                    )
                } else {
                    panic!("Invalid target type")
                }
            })
            .collect();
        let targets = Tensor::cat(targets, 0);

        let images = items
            .into_iter()
            .map(|item| {
                let shape = Shape::new([item.image_height, item.image_width, 3]);
                let data = TensorData::new(image_as_vec_u8(item), shape);
                Tensor::<B, 3>::from_data(data.convert::<B::FloatElem>(), device)
                    .swap_dims(2, 1) // [H, W, C] -> [H, C, W]
                    .swap_dims(1, 0) // [H, C, W] -> [C, H, W]
                    .unsqueeze::<4>() // [C, H, W] -> [1, C, H, W]
                    / 255.0
            }).collect();

        let images = Tensor::cat(images, 0);

        FdsBatch {
            images,
            targets,
        }
    }
}

モデル部分の実装

今回は単純に畳み込み層を二つ持つだけのCNNを定義して回してみます。
まずは実装に必要なものをインポートします。

fds/model.rs

use burn::{
    module::Module,
    nn::{
        LeakyRelu, LeakyReluConfig, Linear, LinearConfig,
        conv::{Conv2d, Conv2dConfig},
        pool::{MaxPool2d, MaxPool2dConfig},
        loss::CrossEntropyLossConfig,
    },
    prelude::*,
    train::{ClassificationOutput, TrainOutput, TrainStep, ValidStep},
    tensor::backend::AutodiffBackend,
};

続いて、Modelの持つ処理層や活性化関数を構造体で定義します。

fds/model.rs

#[derive(Module, Debug)]
pub struct Model<B: Backend> {
    conv2d_1: Conv2d<B>,
    conv2d_2: Conv2d<B>,
    pool: MaxPool2d,
    linear1: Linear<B>,
    linear2: Linear<B>,
    activation: LeakyRelu,
}

次は処理層ごとのパラメータを設定していくんですが、この辺りはPyTorchライクな感じなので、Rust経験がなくてもやってることは理解しやすいかったです。
今回設定したモデルのパラメータは何か設定意図などはなく、テスト用としてとりあえず設定した数字になっています。

fds/model.rs

impl<B: Backend> Model<B> {
    pub fn new(num_classes: usize, device: &B::Device) -> Self {
        let conv2d_1 = Conv2dConfig::new([3, 32], [3, 3]).init(device);
        let conv2d_2 = Conv2dConfig::new([32, 64], [3, 3]).init(device);
        let pool = MaxPool2dConfig::new([2, 2]).init();
        // 224x224が画像サイズ、padding無しなので、(((224-2)/2-2)/2)^2 * 64 = 186624
        // サイズ計算は小数点以下は切り捨てが入ります
        let linear1 = LinearConfig::new(186624, 512).init(device);
        let linear2 = LinearConfig::new(512, num_classes).init(device);
        let activation = LeakyReluConfig::new().init();

        Self {
            conv2d_1,
            conv2d_2,
            pool,
            linear1,
            linear2,
            activation,
        }
    }

    pub fn forward(&self, x: Tensor<B, 4>) -> Tensor<B, 2> {
        let x = self.conv2d_1.forward(x);
        let x = self.activation.forward(x);
        let x = self.pool.forward(x);

        let x = self.conv2d_2.forward(x);
        let x = self.activation.forward(x);
        let x = self.pool.forward(x);

        let x = x.flatten(1, 3);

        let x = self.linear1.forward(x);
        let x = self.activation.forward(x);
        self.linear2.forward(x)
    }

    pub fn forward_classification(
        &self,
        images: Tensor<B, 4>,
        targets: Tensor<B, 1, Int>,
    ) -> ClassificationOutput<B> {
        let output = self.forward(images);
        let loss = CrossEntropyLossConfig::new()
            .init(&output.device())
            .forward(output.clone(), targets.clone());
        ClassificationOutput::new(loss, output, targets)
    }
}

最後にステップごとの処理を実装します。先ほどのモデルに対して、FdsBatchを使ってTrainStepとValidStepトレイトを実装しました。TrainStepの方だけlossから誤差逆伝播させて重みの計算をしています。

fds/model.rs

use crate::fds::data::FdsBatch;

impl<B: AutodiffBackend> TrainStep<FdsBatch<B>, ClassificationOutput<B>> for Model<B> {
    fn step(&self, batch: FdsBatch<B>) -> TrainOutput<ClassificationOutput<B>> {
        let item = self.forward_classification(batch.images, batch.targets);
        TrainOutput::new(self, item.loss.backward(), item)
    }
}

impl<B: Backend> ValidStep<FdsBatch<B>, ClassificationOutput<B>> for Model<B> {
    fn step(&self, batch: FdsBatch<B>) -> ClassificationOutput<B> {
        self.forward_classification(batch.images, batch.targets)
    }
}

学習の振る舞いを実装

学習時にどう処理を進めるかの部分を実装します。
まずはインポートします。

fds/train.rs

use std::fs;
use burn::{
    data::{
        dataloader::DataLoaderBuilder
    },
    optim::AdamConfig,
    prelude::*,
    record::CompactRecorder,
    tensor::backend::AutodiffBackend,
    train::{
        metric::{AccuracyMetric, LossMetric},
        LearnerBuilder,
    },
};
use burn::data::dataset::vision::ImageFolderDataset;

use crate::fds::data::FdsBatcher;
use crate::fds::data::FdsLoader;
use crate::fds::model::Model;

学習用パラメータの設定

Configの構造体を定義し、エポック数やバッチ数などのパラメータを設定します。

fds/train.rs

#[derive(Config, Debug)]
pub struct TrainingConfig {
    pub model: ModelConfig,
    pub optimizer: AdamConfig,
    #[config(default = 10)]
    pub num_epochs: usize,
    #[config(default = 8)]
    pub batch_size: usize,
    #[config(default = 4)]
    pub num_workers: usize,
    #[config(default = 42)]
    pub seed: u64,
    #[config(default = 1.0e-4)]
    pub learning_rate: f64,
}

#[derive(Config, Debug)]
pub struct ModelConfig {}

学習用関数の定義

データローダーを作成し、学習用データとテストデータを読み込み、learnerを作成。
ここまで実装してきたものを使います。こう見るとメソッドチェーンで指定できるとすごく見やすくスッキリして見えるのは自分だけでしょうか？

最後にlearner.fitして学習を行い、モデルを保存して終了です！

fds/train.rs

pub fn train<B: AutodiffBackend>(artifact_dir: &str, config: TrainingConfig, device: B::Device) {
    B::seed(&device, config.seed);

    let batcher = FdsBatcher::new();

    let dataloader_train = DataLoaderBuilder::new(batcher.clone())
        .batch_size(config.batch_size)
        .shuffle(config.seed)
        .num_workers(config.num_workers)
        .build(ImageFolderDataset::fds_train());

    let dataloader_test = DataLoaderBuilder::new(batcher)
        .batch_size(config.batch_size)
        .shuffle(config.seed)
        .num_workers(config.num_workers)
        .build(ImageFolderDataset::fds_test());

    let num_classes = fs::read_dir("data/FDS/train")
        .expect("trainデータディレクトリが見つかりません")
        .filter_map(|entry| entry.ok())
        .filter(|entry| entry.path().is_dir())
        .count();
        
    let model = Model::<B>::new(num_classes, &device);

    let learner = LearnerBuilder::new(artifact_dir)
        .metric_train_numeric(AccuracyMetric::new())
        .metric_valid_numeric(AccuracyMetric::new())
        .metric_train_numeric(LossMetric::new())
        .metric_valid_numeric(LossMetric::new())
        .with_file_checkpointer(CompactRecorder::new())
        .num_epochs(config.num_epochs)
        .summary()
        .build(
            model,
            config.optimizer.init(),
            config.learning_rate,
        );

    let model_trained = learner.fit(dataloader_train, dataloader_test);

    model_trained
        .model
        .save_file(format!("{artifact_dir}/model"), &CompactRecorder::new())
        .expect("モデルの保存に失敗しました");

    println!("🔥 学習完了！モデルを保存しました: {artifact_dir}/model");

}

main.rsを実装

cargo runをした時のエントリーポイントになるコードを実装します。
まずインポートと基本設定です。ここで計算リソースとしてCPUを使うか、GPUを使うかなど選べます。

main.rs

mod fds;

use burn::backend::{
    Autodiff, 
    Cuda, 
    // Wgpu
};
use burn::optim::AdamConfig;

use fds::train::{ModelConfig, TrainingConfig};

type MyBackend = Cuda<f32, i32>;
// type MyBackend = Wgpu<f32, i32>;

type MyAutodiffBackend = Autodiff<MyBackend>;

GPUバックエンド選び

DGX Sparkで今回試したのは以下の二つです。

• Cuda：NVIDIA GPU向け（CUDAを利用するバックエンド）
• Wgpu：wgpu（WebGPU）経由で Vulkan/Metal/DX12 などに載せるバックエンド

使った感じとして、軽いタスク（MNISTなど）ならwgpuでもDGX Spark上では動きます。しかし、まだCudaの方が、速くて、動作が安定しているのでCudaバックエンドを使った方がいいと思います。

wgpuでは、単純にデータサイズが大きかったり、層が深いモデルを使おうとすると、エラーで停止しがちでした。
使う際は以下のような感じで簡単に切り替えが可能です。

use burn::backend::Cuda;
type MyBackend = Cuda<f32, i32>;

use burn::backend::Wgpu;
type MyBackend = Wgpu<f32, i32>;

mainを実装

デバイスは上記の方法で定義しておくと、あとはDefault::default()を呼ぶことでその設定を取得できるようになります。なのでtype MyBackend = 〇〇の部分を変更するだけでバックエンドを変更できます（超便利）。

あとは学習用パラメーターなどを設定し、trainを実行すれば学習されるという流れになります。

main.rs

fn main() {
    let device = Default::default();


    // 学習設定を作成
    let config = TrainingConfig::new(
        ModelConfig::new(),
        AdamConfig::new(),
    )
        .with_num_epochs(5)
        .with_batch_size(16)
        .with_learning_rate(1e-3)
        .with_num_workers(1);

    // 学習を実行
    fds::train::train::<MyAutodiffBackend>(
        "./artifacts",
        config,
        device,
    );

}

学習実行

TUIで確認

個人的にBurnの推し機能で、学習曲線や現在のプログレスなどをTUI表示してくれます！学習中の様子がとても見やすいです。ただ、重い学習をさせると表示がされなくなったりもします。

DGX Dashboardで使用率を見る

使用率などはVRAMじゃなくてSystem Memoryになっていたりするため、素直にDashboardを見た方が見やすかったです。

学習結果

最後にLearning Summaryとしてこのように出力してくれます。結構いい感じに分類できているんじゃないでしょうか！

DGX Sparkの使用感について

実際にディープラーニングで使ってみて思ったこととして、以下のような感想を持ちました。

• メモリ数が多く大きなデータや大きいモデルを載せやすい
• GPU自体の処理速度がそこまで速くないため、学習速度は向上しにくい

これは実際にインターン生から聞いた話とも合致していました。以前はGCP上のColab Enterpriseを使っていたんですが、メモリ容量が多くなったことでバッチサイズを増やしたり重いデータを載せることができるようになった一方で、速度はそんなに変わらないという感じでした。

とはいえ、学習速度は早ければいいですけど、学習中って基本的に寝かせておくことが多いのでそこまで気にしなくていいんじゃないかな？とも思います。大事なのは大きいデータと、大きいモデルが乗ることだと思います。

まとめ

今回はDGX SparkでRustのディープラーニングを試してみるということで、Rustで機械学習コードを書く良さを感じつつ、DGX Sparkの使用感がわかってとても良い試みになったかなと思います。
この知見が何かしらの参考となればと思います。

おまけ

今回、fdsというフォルダを作成し、そこに実装していました。
以下のような構成にすることで使用データごとに実装を切り分けておけるのも便利だなと感じました。

src
├── main.rs
├── fds
│   ├── data.rs
│   ├── model.rs
│   └── train.rs
├── fds.rs
├── mnist
│   ├── data.rs
│   ├── model.rs
│   └── train.rs
└── mnist.rs