Burn Framework (第2部): モデル定義とマルチバックエンド構成

Burn ↗ の大きな特長の一つは、モデルアーキテクチャの定義とハードウェア計算バックエンドが完全に分離されている点です。

`Module` マクロによるニューラルネットワークの定義#

Burnでは、すべてのモデル要素が Module トレイトを実装します。#[derive(Module)] プロシージャルマクロにより、レイヤー定義を型安全かつ簡潔に行えます。

src/model.rs

use burn::nn::{Linear, LinearConfig, Relu};
use burn::module::Module;
use burn::tensor::backend::Backend;
use burn::tensor::Tensor;

#[derive(Module, Debug)]
pub struct Model<B: Backend> {
    linear1: Linear<B>,
    linear2: Linear<B>,
    activation: Relu,
}

impl<B: Backend> Model<B> {
    pub fn new(input_dim: usize, hidden_dim: usize, output_dim: usize, device: &B::Device) -> Self {
        let linear1 = LinearConfig::new(input_dim, hidden_dim).init(device);
        let linear2 = LinearConfig::new(hidden_dim, output_dim).init(device);
        Self {
            linear1,
            linear2,
            activation: Relu::new(),
        }
    }

    pub fn forward(&self, input: Tensor<B, 2>) -> Tensor<B, 2> {
        let x = self.linear1.forward(input);
        let x = self.activation.forward(x);
        // [!code focus]
        self.linear2.forward(x)
    }
}

rust

マルチバックエンドの強み#

ジェネリクス型パラメータ B を変更するだけで、実行環境を簡単に切り替えられます：

WGPU Backend: Vulkan, Metal, DirectXに対応した高性能クロスプラットフォームGPU実行。
LibTorch Backend: PyTorch C++ バインディングとの互換性。
Candle Backend: Hugging Face Candleをベースにした軽量バックエンド。
NdArray Backend: CPU専用のフォールバック。

src/main.rs

use burn::backend::wgpu::{Wgpu, WgpuDevice};

fn main() {
    let device = WgpuDevice::default();
    let model = Model::<Wgpu>::new(784, 128, 10, &device);
    println!("モデルの初期化が完了しました！");
}

rust

Module マクロによるニューラルネットワークの定義#

マルチバックエンドの強み#

参考文献#

`Module` マクロによるニューラルネットワークの定義#