ndarray-conv

ndarray-conv 是一个提供纯 Rust 编写的 N 维卷积（带 FFT 加速）库的 crate。

受以下项目启发

ndarray-vision (https://github.com/rust-cv/ndarray-vision)

convolutions-rs (https://github.com/Conzel/convolutions-rs#readme)

pocketfft (https://github.com/mreineck/pocketfft)

路线图

• N 维基本卷积 Array/ArrayView
• N 维卷积带 FFT 加速 Array/ArrayView
• 实现 ConvMode 和 PaddingMode
  • ConvMode: 全部 相同 有效 自定义 明确
  • PaddingMode: 零 常数 反射 复制 循环 自定义 明确
• 带步长的卷积
• 带膨胀的核
• 处理输入大小错误
• 明确的错误类型
• 与类似库进行基准测试

示例

use ndarray_conv::*;

x_nd.conv(
    &k_n,
    PaddingSize::Full,
    PaddingMode::Circular,
);

x_1d.view().conv_fft(
    &k_1d,
    ConvMode::Same,
    PaddingMode::Explicit([[BorderType::Replicate, BorderType::Reflect]]),
);

x_2d.conv_fft(
    k_2d.with_dilation(2),
    PaddingSize::Same,
    PaddingMode::Custom([BorderType::Reflect, BorderType::Circular]),
);

// avoid loss of accuracy for fft ver
// convert Integer to Float before caculate.
x_3d.map(|&x| x as f32)
    .conv_fft(
        &kernel.map(|&x| x as f32),
        ConvMode::Same,
        PaddingMode::Zeros,
    )
    .unwrap()
    .map(|x| x.round() as i32);

fn main() {
    use ndarray_conv::*;
    use ndarray::prelude::*;
    use ndarray_rand::rand_distr::Uniform;
    use ndarray_rand::RandomExt;
    use std::time::Instant;

    let mut small_duration = 0u128;
    let test_cycles_small = 1;
    // small input data
    for _ in 0..test_cycles_small {
        let x = Array::random((2000, 4000), Uniform::new(0., 1.));
        let k = Array::random((9, 9), Uniform::new(0., 1.));

        let now = Instant::now();
        // or use x.conv_fft() for large input data
        x.conv(
            &k,
            ConvMode::Same,
            PaddingMode::Custom([BorderType::Reflect, BorderType::Circular]),
        );
        small_duration += now.elapsed().as_nanos();
    }

    println!(
        "Time for small arrays, {} iterations: {} milliseconds",
        test_cycles_small,
        small_duration / 1_000_000
    );
}

版本

• 0.3.3 - 修复错误：纠正 conv_fft 的输出形状。
• 0.3.2 - 通过修改 good_fft_size 和 transpose 来提高性能。
• 0.3.1 - 实现 basic error 类型。修复一些错误。
• 0.3.0 - 更新为 N 维卷积。
• 0.2.0 - 完成 conv_2d 和 conv_2d_fft。