Rust+Burnで勾配を取得する

BurnはRust製のDLフレームワークです。

burn.dev

実行前に少しだけ次元のチェックをしてくれます。

RustのDeep LearningライブラリBurn、書いている途中からある程度の次元チェックしてくれるのうれしい。 pic.twitter.com/OHS5RIc1IN

— natsutan (@natsutan) 2024年12月14日

PyTorchでいうとdimのチェックをしてくれます。 [1, 3, 64, 64]と[3, 64, 64]は違うとエラーを出してくれますが、[1, 3, 64, 64]と[1, 3, 32, 32]は実行時のチェックになります。

Burnを使って単に自動微分の勾配を取得したかったのですが、それだけをやっているサンプルが見つかりませんでした。 なんとかここまで来れたのでメモ残しです。

ソースコード

rustソース

use burn::tensor::Tensor;  
use burn::backend::{Autodiff, Wgpu};  
  
type Backend = Wgpu;  
type AutoDIffBackend = Autodiff<Backend>;  
  
fn main() {  
    let device = Default::default();  
  
    let x = Tensor::<AutoDIffBackend, 1>::from_floats([5.0], &device).require_grad();  
    let y = x.clone().powf_scalar(2.0) * 3.0;  
  
    // 勾配を求めて表示する  
    let gradients = y.backward();  
    let tensor_grad = x.grad(&gradients).unwrap();  
  
    println!("y = {}", y);  
    println!("dy/dx = {:?}",  tensor_grad.into_scalar());  
  
}

carog.toml

[package]  
name = "ch5_nn"  
version = "0.1.0"  
edition = "2021"  
  
[dependencies]  
burn = { version = "~0.15", features = ["train", "wgpu", "vision", "candle"] }

説明

BackendにはWgpuを使い、さらにAutoDIffBackendを定義する。

type Backend = Wgpu;  
type AutoDIffBackend = Autodiff<Backend>;  

勾配が欲しいTensorにrequire_grad()をつける。 Tensorの所有権が移ってしまうと後で勾配が求められないのでclone()が必要。

    let x = Tensor::<AutoDIffBackend, 1>::from_floats([5.0], &device).require_grad();  
    let y = x.clone().powf_scalar(2.0) * 3.0;  

backward()した後勾配を求める。数値として取り出すには、into_scalar()が必要。

    let gradients = y.backward();  
    let tensor_grad = x.grad(&gradients).unwrap();  
    println!("dy/dx = {:?}",  tensor_grad.into_scalar());  

これで実行するとちゃんと勾配が求められています。

y = Tensor {
  data:
[74.99999],
  shape:  [1],
  device:  BestAvailable,
  backend:  "autodiff<fusion<jit<wgpu<wgsl>>>>",
  kind:  "Float",
  dtype:  "f32",
}
dy/dx = 30.0

ちなみにBurnはWSL2で動かすと遅いですが、Windowsで直接動かすと爆速です。 Windows環境の人はWindowsで動かす事をお勧めです。