Rustの機械学習ライブラリCandleで果物判定を実装してみよう

• 

  Rustの機械学習ライブラリCandleを使って果物判定をしよう

機械学習ライブラリ「Candle」とは

Candleは、Hugging Faceが開発するRust製の機械学習フレームワークです。Rustの安全性と高速性を活かし、Python環境に依存せずモデルの推論や学習を行える点が特徴です。

• 

  CandleのGitHubリポジトリ

Candleの構成はシンプルで、candle-coreによるテンソル演算、candle-nnによるニューラルネットワーク層、candle-optimisersによる最適化器などに分かれています。PyTorch風のAPIを備えつつRustらしいモジュール設計がなされており、CPUだけでなくCUDAやMPSなどGPUも利用できるのも魅力です。

また、生成AIのモデルを多数ホストしているHugging Faceが開発しているだけあって、Hugging Face Hub上の大規模言語モデル(LLM)や画像生成モデルをRustから直接扱えるようになっています。そのため、組み込みやCLIツール、WebAssemblyを活用したブラウザ上での推論なども可能です。

Candleの実力を果物判定で確かめよう

ここでは、Candleの実力を手軽に確かめるために、クリエイティブコモンズのライセンスの元に公開されている果物画像のデータセット「fruits_db」を活用してみましょう。このデータセットには、たくさんの果物画像が含まれており、機械学習を利用して、レモンの画像とイチゴの画像を判定するタスクをテストすることができます。

• 

  果物画像のデータセット「fruits_db」

まずは、果物データセットを、こちらのダウンロードページからZIP形式でダウンロードしましょう。ダウンロードして解凍すると、次のようなフォルダ構造となっています。

.
├── data
│   ├── lemon
│   │   ├── 9474250579.jpg --- レモンの画像
│   │   ├── 9474255143.jpg
│   │   ├── 9628093234.jpg
│   │   ├──  …
│   └── strawberry
│       ├── 11074319503.jpg --- イチゴの画像
│       ├── 112644688.jpg
│       ├── 11295203274.jpg
│       ├──  …
└── README.md

このように、果物のJPEG画像が、果物ごとにフォルダ分けされて収録されたものとなっています。

• 

  fruits_dbは果物ごとに画像ファイルが収録されたもの

Rustの果物分類プロジェクトを作成しよう

まずは、Rustのプロジェクトを作成しましょう。ターミナル(WindowsならPowerShell、macOSならターミナル.app)で以下のコマンドを実行して、プロジェクトを初期化しましょう。

# プロジェクトのフォルダを作成
mkdir program
cd program
# プロジェクトを初期化
cargo init
# 必要なライブラリを追加
cargo add anyhow rand walkdir image@0.25
cargo add candle-core@0.9 candle-nn@0.9 candle-optimisers@0.9    

そして、program以下に、果物データセットのdataフォルダをコピーしましょう。次のようなフォルダ構成にします。

.
├── Cargo.lock
├── Cargo.toml
├── data --- fruits_dbのdataフォルダをここにコピー
│   ├── lemon
│   └── strawberry
├── src
│   └── main.rs --- メインプログラム
└── target

プログラムを実装しよう

それでは、プログラムを実行しましょう。ここで作成したプログラムは、次のようなもので222行あります。

• 

  Candleで深層学習モデルのMLPを実装したもの

機械学習のプログラムにしては短いのですが、連載で紹介するには、少し長いものとなっています。そこで、こちらにコード全部をアップロードしました。コードをコピーして、src/main.rs に貼り付けて保存しましょう。

プログラムを実行しよう

プログラムを実行するには、ターミナルで下記のコマンドを実行しましょう。

cargo run

プログラムが正しく実行できると、次のように表示されます。

• 

  分類精度を確認すると、98.46%を達成したところ

機械学習では、学習用データとテスト用データにデータセットを分割し、学習用データを使って分類モデルを作成し、テスト用データを使って、どのくらいの精度が出るかをテストします。

上記の実行結果を見ると、最終的に98.46%の精度を達成することができたことを表しています。ただし、毎回データをシャッフルしてから学習するため、最終的な分類精度が90%前後の場合もあります。

プログラムのポイントを確認しよう

プログラムが正しく実行できたら、簡単にプログラムのポイントを確認しておきましょう。

以下のコードは、機械学習のパラメータを定義している部分です。定数IMG_SIZEは、画像サイズを32x32にリサイズしてから学習するようにします。ここでは、画像サイズ×3色のRGB画像として扱うため、入力次元は3072となります。また、学習時のエポック数やバッチサイズも定義しています。

// 機械学習で使うパラメータを定義 --- (*1)
const IMG_SIZE: u32 = 32; // 画像サイズ(32, 32) * 3色=3072次元
const EPOCHS: usize = 20; // 学習時のエポック数
const BATCH_SIZE: usize = 128; // 学習時のバッチサイズ

以下の部分では、データセットの1サンプルを表現する構造体Sampleを定義しています。

/// データセットの1サンプルを表現する構造体 --- (*2)
#[derive(Clone)]
struct Sample {
    x: Vec<f32>, // 画像データ(32x32x3=3072次元)
    y: u32,      // クラスID
}

そして、以下の部分がMLP(Multi Layer Perceptron)のモデルを定義している部分です。全結合層を持つ3層のニューラルネットワークを構成しています。

/// 単純なMLP構造のモデルを定義 --- (*3)
struct Mlp {
    // 全結合層（candle_nn::Linear）を3段積んだ 多層パーセプトロン
    l1: candle_nn::Linear,
    l2: candle_nn::Linear,
    l3: candle_nn::Linear,
}
impl Mlp {
    /// MLPモデルを構築する(vbは重み・バイアス, in_dimは入力次元, out_dimは出力次元)
    fn new(vb: VarBuilder, in_dim: usize, hidden1: usize, hidden2: usize, out_dim: usize) -> Result<Self> {
        Ok(Self {
            l1: linear(in_dim, hidden1, vb.pp("l1"))?,
            l2: linear(hidden1, hidden2, vb.pp("l2"))?,
            l3: linear(hidden2, out_dim, vb.pp("l3"))?,
        })
    }
}

コードは掲載しませんが、プログラムの(＊4)の部分で、dataフォルダを走査して、果物画像を読み込み、32x32にリサイズして、Sample構造体のベクタに格納しています。そして、(＊5)の部分で画像データセットを読み込みます。(＊6)でデータセットをシャッフルして、学習用データとテスト用データに分割しています。

以下のコード(＊7)の部分で、Sample構造体のベクタをCandleのTensorに変換します。

// データをTensor化する --- (*7)
let device = Device::Cpu;
let (x_train, y_train) = to_tensors(train, in_dim, &device)?;
let (x_test, y_test) = to_tensors(test, in_dim, &device)?;
…省略…
/// SampleのスライスをTensorに変換
fn to_tensors(samples: &[Sample], in_dim: usize, device: &Device) -> Result<(Tensor, Tensor)> {
    let mut xs = Vec::with_capacity(samples.len() * in_dim);
    let mut ys = Vec::with_capacity(samples.len());
    for s in samples {
        anyhow::ensure!(s.x.len() == in_dim, "unexpected dim {}", s.x.len());
        xs.extend_from_slice(&s.x);
        ys.push(s.y);
    }
    let x = Tensor::from_vec(xs, (samples.len(), in_dim), device)?;
    let y = Tensor::from_vec(ys, samples.len(), device)?.to_dtype(DType::U32)?;
    Ok((x, y))
}

そして、(＊8)の部分でMLPモデルを構築し、以下のコード(＊9)の部分でデータの学習を行います。(＊10)の部分で、テスト用データを使って、学習したモデルの精度を確認します。

// 学習ループ --- (*9)
let n_train = x_train.dims()[0];
println!("train: {n_train}, test: {}", x_test.dims()[0]);
for epoch in 1..=EPOCHS {
    let mut idx: Vec<usize> = (0..n_train).collect();
    idx.shuffle(&mut rng());
    let mut total = 0f32;
    let mut steps = 0usize;

    for chunk in idx.chunks(BATCH_SIZE) {
        let chunk_u32: Vec<u32> = chunk.iter().map(|&x| x as u32).collect();
        let chunk_tensor = Tensor::from_vec(chunk_u32, chunk.len(), &device)?;
        let xb = x_train.index_select(&chunk_tensor, 0)?;
        let yb = y_train.index_select(&chunk_tensor, 0)?;
        let logits = model.forward(&xb)?;
        let loss = loss::cross_entropy(&logits, &yb)?;
        total += loss.to_scalar::<f32>()?;
        steps += 1;
        opt.backward_step(&loss)?;
    }
    // テスト精度を確認 --- (*10)
    let acc = accuracy(&model, &x_test, &y_test)?;
    println!("epoch {epoch}: loss={:.4}, acc={:.2}%", total / steps as f32, acc * 100.0);
}

macOSのMシリーズのGPU(Metal)に対応するよう改良しよう

ところで、Candleは、GPUにも対応しています。プログラムをちょっと修正することで、GPU対応させることが可能です。以下は、macOS (Apple Silicon) の、Metal Performance Shaders (MPS) を使うように修正する方法です。

最初に、Cargo.tomlで、MPSが有効になるように、featuresを指定します。修正したら「cargo clean」を実行してキャッシュを削除しましょう。

candle-core = { version = "0.9", features = ["metal"] }
candle-nn = { version = "0.9", features = ["metal"] }

そして、プログラムの(*7)のところを下記のように修正します。

// データをTensor化する --- (*7)
let device = match Device::new_metal(0) {
    Ok(device) => { // MPSの特定に成功した時
        println!("Using Metal GPU acceleration");
        device
    }
    Err(_) => { // MPS非搭載の機種の場合
        println!("Metal not available, falling back to CPU");
        Device::Cpu
    }
};

筆者の手元にあるMacBook Pro M4(メモリ32GB)で試してみましたが、CPUを利用する場合は、プログラムの実行に5.484秒かかったのですが、GPUを利用するように修正したものは、1.83秒で学習することができました。約3倍の速度で実行させることができました。

まとめ

以上、今回は、Candleを利用して果物画像の分類を行うプログラムを実装してみました。Candleを使えば、Pythonに依存せず、ローカル環境で深層学習モデルを動作させることができるのが魅力です。今回のプログラムは単純なMLPを実装したものではありますが、macOSのバイナリは4.8MB程度のサイズでした。ぜひ、皆さんもCandleを活用して、いろいろな機械学習タスクに挑戦してみてください。

自由型プログラマー。くじらはんどにて、プログラミングの楽しさを伝える活動をしている。代表作に、日本語プログラミング言語「なでしこ」 、テキスト音楽「サクラ」など。2001年オンラインソフト大賞入賞、2004年度未踏ユース スーパークリエータ認定、2010年 OSS貢献者章受賞。これまで50冊以上の技術書を執筆した。直近では、「大規模言語モデルを使いこなすためのプロンプトエンジニアリングの教科書(マイナビ出版)」「Pythonでつくるデスクトップアプリ(ソシム)」「実践力を身につける Pythonの教科書 第2版」「シゴトがはかどる Python自動処理の教科書(マイナビ出版)」など。