はじめに

画像処理などでループ処理を統合してメモリアクセス回数を減らすことで、
性能を向上させることができる場合があります。

これを Rust のイテレータで書くと、for文で書くより嬉しいことが
起こらないかというのが、今回実験してみた内容です。

まずは Python で書いてみる

例えば Python で numpy を使って下記のように書くような処理を考えます。

import numpy as np
import time

WIDTH = 8192
HEIGHT = 4096

# ランダムな値を画像の代わりに使う
a = np.random.random((HEIGHT, WIDTH)) # 画像1
b = np.random.random((HEIGHT, WIDTH)) # 画像2

start_time = time.time() # 時間計測開始
rmse = np.sqrt(np.sum((a - b)**2)/(HEIGHT*WIDTH)) # 演算してみる
duration = (time.time() - start_time) * 1000 # 時間計測終了
print(sum)
print(f"time:{duration:5.2f}[ms]")

ちなみに私の環境での実行時間は

time:224.41[ms]

でした。

この画像処理演算の内容としては

２つの画像の差分を取って誤差画像を生成する処理
誤差画像を二乗する処理
二乗した誤差画像の総和を取る処理
総和を画像サイズで割って二乗平均を得る処理
二乗平均の平方根を計算して二乗平均平方根誤差を求める処理

に分解され、numpy の ndarray の単位で、それぞれで処理が行われます。

ちなみに C++ や Rust で、OpenCV などを使う場合も cv::Mat 型の単位で演算が行われるので、同じことが起こります。

そのまま Rust に置き換えてみる

上の numpy の計算順序のままになるように Rust のコードに展開すると

use std::time;
use rand::prelude::*;

const WIDTH  : usize = 8192;
const HEIGHT : usize = 4096;

fn main() {
    // ランダムな値を画像の代わりに使う
    let mut a : Vec<[f64; WIDTH]> = vec![[0.0; WIDTH]; HEIGHT];
    let mut b : Vec<[f64; WIDTH]> = vec![[0.0; WIDTH]; HEIGHT];
    let mut rng    = rand::thread_rng();
    for y in  0..HEIGHT {
        for x in  0..WIDTH {
            a[y][x] = rng.gen();
            b[y][x] = rng.gen();
        }
    }

    // 作業領域
    let mut c : Vec<[f64; WIDTH]> = vec![[0.0; WIDTH]; HEIGHT];
    let mut d : Vec<[f64; WIDTH]> = vec![[0.0; WIDTH]; HEIGHT];

    // 時間計測開始
    let start_time = time::Instant::now();

    // ２つの画像の差分を取って誤差画像を生成する処理
    for y in  0..HEIGHT {
        for x in  0..WIDTH {
            c[y][x] = a[y][x] - b[y][x];
        }
    }

    // 誤差画像を二乗する処理
    for y in  0..HEIGHT {
        for x in  0..WIDTH {
            d[y][x] = c[y][x] * c[y][x];
        }
    }
    
    // 二乗した誤差画像の総和を取る処理
    let mut sum : f64 = 0.0;
    for y in  0..HEIGHT {
        for x in  0..WIDTH {
            sum += d[y][x];
        }
    }

    // 総和を画像サイズで割って二乗平均を得る処理
    let mean = sum / (HEIGHT*WIDTH) as f64;

    // 二乗平均の平方根を計算して二乗平均平方根誤差を求める処理
    let msre = mean.sqrt();
    
    // 時間計測終了
    let duration  = start_time.elapsed().as_secs_f64() * 1000.0;

    println!("sum = {}", msre);
    println!("time:{:5.2}[ms]", duration);
}

というようなコードになり、Release ビルド時の実行速度は

time:136.74[ms]

でした。

ループを統合してみる

これは、複数ある for 文を統合して計算部分を下記のように書き直すことで高速化します。

    // 計算を1つのループにまとめる
    let mut sum : f64 = 0.0;
    for y in  0..HEIGHT {
        for x in  0..WIDTH {
            let c = a[y][x] - b[y][x];
            let d = c * c;
            sum += d;
        }
    }
    let mean = sum / (HEIGHT*WIDTH) as f64;
    let msre = mean.sqrt();

作業用のメモリが不要となり、メモリアクセス量が減るとともにキャッシュヒット率も上がることが期待できます。

私の環境で

time:34.99[ms]

となり、かなり高速化しました。

イテレータで書いてみる

しかしながら、元の python のコードと比べるとだいぶ見通しの悪いコードになってしまいます。

そこで、 Rust のイテレータを試してみます。イテレータあまり慣れてないので GPT-4 に助けてもらいました。

    // 計算をイテレータで書く
    let msre : f64 = (a.iter().zip(b.iter())
        .flat_map(|(row1, row2)| row1.iter().zip(row2.iter()))
        .map(|(a, b)| a - b)
        .map(|c| c*c)
        .sum::<f64>()
        / (HEIGHT*WIDTH) as f64).sqrt();

実行時間は for 文を使った場合とおおむね同じで

time:35.92[ms]

となりました。

まとめ

今回は単純な演算だったのであまりメリットが分かりずらいですが、複数の画像処理フィルタを組み合わせて処理するような用途だと、各処理をイテレータ―として利用できるように個別に作成しておけば、「自由に組み換えができる」というメリットと「性能を出す」というメリットの両方が目指せる可能性があるように感じました。

Discussion