映画と時々技術の日記

データ分析やNLPが専門です.よく映画を見ます.

Rustで数理最適化を試してみる

最近Rustの公式ドキュメントとrustlingsをちょいちょいやっていたので,ここらで一度何かしらをしてみようと思い数理最適化をやる話です.具体的にはargminを使ってドキュメントにあるexampleを動かしてみます.

argminとは

Rustで書かれている数理最適化のツールボックス/フレームワークです.ドキュメントを見た感じ,Pulpやgurobipyなどとはまた違った書き方をしていて,少し面白そうなのでこちらを使ってみます.ちなみに,Googlerust numerical optimization で検索をすると最初にOptimizeというcrateが出てきます.Optimizeはscipy.optimizeをベースにして作成されているので,SciPyをよく使っている方にはなじみ深いのかなと思います.

やってみる

環境はUbuntu 20.04 LTS,rustcのバージョンは1.55.0です.exampleを見るに,Cargo.tomlに必要なcrateを追加して,例として出ているコードを実行したら動きそうです.

最初に,Cargoを使って適当な環境を作成します.

cargo new argmin_test

次に,Cargo.tomlにargminを追加します.ドキュメントでは2通りの方法が書かれていますが,とりあえず推奨で行います.ドキュメントではwindowsの場合に ndarray-linalg を追加することなど書いてありますが,今回はスルーします.

[dependencies]
argmin = { version = "0.4.7", features = ["ctrlc", "ndarrayl", "nalgebral"] }

ドキュメントにあるサンプルのコードをまるごとコピペします.Rosenbrock という構造体を定義し,その中に最適化に用いる定数を定義します.機械学習のパラメータの最適化などであれば,教師データをこちらに放り込むようなイメージです.そして定義した構造体に対してArgminOpを実装することでモデルを定義します.

use argmin_testfunctions::{rosenbrock_2d, rosenbrock_2d_derivative, rosenbrock_2d_hessian};
use argmin::prelude::*;

/// First, create a struct for your problem
struct Rosenbrock {
    a: f64,
    b: f64,
}

/// Implement `ArgminOp` for `Rosenbrock`
impl ArgminOp for Rosenbrock {
    /// Type of the parameter vector
    type Param = Vec<f64>;
    /// Type of the return value computed by the cost function
    type Output = f64;
    /// Type of the Hessian. Can be `()` if not needed.
    type Hessian = Vec<Vec<f64>>;
    /// Type of the Jacobian. Can be `()` if not needed.
    type Jacobian = ();
    /// Floating point precision
    type Float = f64;

    /// Apply the cost function to a parameter `p`
    fn apply(&self, p: &Self::Param) -> Result<Self::Output, Error> {
        Ok(rosenbrock_2d(p, self.a, self.b))
    }

    /// Compute the gradient at parameter `p`.
    fn gradient(&self, p: &Self::Param) -> Result<Self::Param, Error> {
        Ok(rosenbrock_2d_derivative(p, self.a, self.b))
    }

    /// Compute the Hessian at parameter `p`.
    fn hessian(&self, p: &Self::Param) -> Result<Self::Hessian, Error> {
        let t = rosenbrock_2d_hessian(p, self.a, self.b);
        Ok(vec![vec![t[0], t[1]], vec![t[2], t[3]]])
    }
}

それでは実際にソルバを動かしてみます.ドキュメントにあるexampleを貼り付け,コマンドラインからcargo runで実行します.

use argmin::prelude::*;
use argmin::solver::gradientdescent::SteepestDescent;
use argmin::solver::linesearch::MoreThuenteLineSearch;

// Define cost function (must implement `ArgminOperator`)
let cost = Rosenbrock { a: 1.0, b: 100.0 };
  
// Define initial parameter vector
let init_param: Vec<f64> = vec![-1.2, 1.0];
  
// Set up line search
let linesearch = MoreThuenteLineSearch::new();
  
// Set up solver
let solver = SteepestDescent::new(linesearch);
  
// Run solver
let res = Executor::new(cost, solver, init_param)
    // Add an observer which will log all iterations to the terminal
    .add_observer(ArgminSlogLogger::term(), ObserverMode::Always)
    // Set maximum iterations to 10
    .max_iters(10)
    // run the solver on the defined problem
    .run()?;
  
// print result
println!("{}", res);

ドキュメントのexampleをほぼ脳死でコピペして実行すると,2つのエラーにぶち当たりました.

  1. 使っているcrateがない
  2. ?オペレーターの有無

最初に,crateがないものに関してです.以下のようなエラーが出てきました.

error[E0432]: unresolved import `argmin_testfunctions`
 --> src/main.rs:1:5
  |
1 | use argmin_testfunctions::{rosenbrock_2d, rosenbrock_2d_derivative, rosenbrock_2d_hessian};
  |     ^^^^^^^^^^^^^^^^^^^^ use of undeclared crate or module `argmin_testfunctions`

モデルの定義の場所で使っているargmin_testfunctionsというのが存在しないというエラーです.Cargo.tomlを見ても,確かにそんなものを入れていないので当然のエラーです.というわけで探してきて以下を追記します.

argmin_testfunctions = "0.1.1"

2番目に出てきたのが?オペレーターの有無に関してです.

error[E0277]: the `?` operator can only be used in a function that returns `Result` or `Option` (or another type that implements `FromResidual`)
   --> src/main.rs:62:15
    |
42  | / fn main() {
43  | |     // Define cost function (must implement `ArgminOperator`)
44  | |     let cost = Rosenbrock { a: 1.0, b: 100.0 };
45  | |
...   |
62  | |         .run()?;
    | |               ^ cannot use the `?` operator in a function that returns `()`
...   |
65  | |     println!("{}", res);
66  | | }
    | |_- this function should return `Result` or `Option` to accept `?`

まだほとんど?オペレーターに関して理解していないのでなんで出るんだろうなぁと思いながらとりあえずエラーを確認します.見た感じrun()の返り値がResultOptionじゃないと?オペレーターはつけることができないと書いてあります.正直,run()の返り値はResultで普通に返ってきてそうなんだけどなぁ…と思っています.とりあえず動かすことを優先して私がとった対策としては,?オペレーターが使えないならunwrap()を後ろにつけて結果をとってくるようにしました.

ここまでやってとりあえず動く形にはなりました.ターミナルの方では最適化の経過ログが吐かれ,最終的な結果を出力すると以下のようになっています.

ArgminResult:
    param (best):  [-1.0094914782615345, 1.0268328735150087]
    cost (best):   4.044077495556465
    iters (best):  9
    iters (total): 10
    termination: Maximum number of iterations reached
    time:        Some(7.365ms)

おわりに

とりあえず,argminのexampleを動かすことに成功しました.

他にも最適化中のログをいじったり,最適化アルゴリズムを自分で定義できたりとかなり自由度が高そうだと思います.使いこなせるかわからないですがいろいろ試していこうかなと思います.ここまで勢いでそのまま書いてきたのですが,Pulpなんかとの比較のコードがあったりするほうが備忘録としても良さそうなので元気が余っていればそちらもやっていきたいですね.