Kaffe - 一个受 Pytorch 启发的 Rust 库

Kaffe 如标题所示，是一种在 Rust 中创建神经网络的方法。

目标是创建一种简单的方式来编写自己的模型并测试它们。语法应该与 Pytorch 类似，但一些功能可能会借鉴 numpy 或甚至 tensorflow 的名称。

将来，矩阵库可能会完全转移到自己的项目中，但现在它们都在同一个 crate 中；

为什么？因为有时候你想要在 Rust 中编写酷炫且快速的东西 :)

示例

矩阵基本示例

use kaffe::tensor::Tensor;

fn main() {
    let t = Tensor::init(10f32, vec![2, 2, 2]);

    let res = t.log(10.0);

    println!("{:?}", res.data);

    let tensor = Tensor::randomize_range(1.0, 4.0, vec![2, 4]);

    assert_eq!(tensor.all(|&e| e >= 1.0), true);

    let tensor = Tensor::init(20.0, vec![2, 2]);
    let value: f32 = 2.0;

    let result_mat = tensor.div_val(value);

    assert_eq!(result_mat.data, vec![10.0; 4]);

    let tensor = Tensor::init(4f32, vec![1, 1, 1, 4]);

    assert_eq!(tensor.data, vec![4f32; 4]);
    assert_eq!(tensor.shape, vec![1, 1, 1, 4]);

    let mut tensor = Tensor::init(2.0, vec![2, 4]);
    println!("{}", tensor.data[0]);

    tensor.set_where(|e| {
        if *e == 2.0 {
            *e = 2.3;
        }
    });

    println!("{}", tensor.data[0]);

    assert_eq!(tensor.data[0], 2.3);

    println!("{}", tensor.get(vec![0, 0]).unwrap());
}

神经网络基本示例 - 待实现

use kaffe::Matrix;
use kaffe::{Net, Layer, optimizer::*, loss::*};

// Here lies our model 
struct MyNet {
    layers: Vec<Layer>
}

// Implement default functions
impl Net for MyNet {
    /// Set's up parametes for the struct 
    fn init() -> Self {
        let mut layers: Vec<Layers> = Vec::new();
        self.layers.push(nn.Conv2d(1,32,3,1));
        self.layers.push(nn.Conv2d(32,64,3,1));
        self.layers.push(nn.Dropout(0.25));
        self.layers.push(nn.Dropout(0.5));
        self.layers.push(nn.FCL(9216, 128));
        self.layers.push(nn.FCL(128,10));

        Self { layers }
    }

    /// Define a forward pass 
    fn forward(x: &Matrix) {
        x = layers[0](x)
        x = ReLU(x);
        x = layers[1](x)
        x = ReLU(x);
        x = layers[2](x)
        x = ReLU(x);
        let output = log_softmax(x);
        return output;
    }
}

fn train(model: &Model, 
        train_dataloader: &DataLoader, 
        optimizer: &Optimizer, 
        epoch: usize) {
    model.train();

    for (batch_idx, (data, target)) in train_dataloader.iter().enumerate() {
        optimizer.zero_grad();
        let output = model(data);
        let loss = BCELoss(output, target);
        loss.backward();
        optimizer.step();
    }
}

fn test(model: &Model, 
        test_dataloader: &DataLoader, 
        optimizer: &Optimizer, 
        epoch: usize) {
    model.eval();

    let mut test_loss = 0.0;
    let mut correct = 0.0;

    optimizer.no_grad();

    for (batch_idx, (data, target)) in train_dataloader.iter().enumerate() {
        let output = model(data);
        test_loss += BCELoss(output, target);

        let pred = output.argmax(Dimension::Row);
        correct += pred.eq(target.view_as(pred)).sum();
    }
    test_loss /= test_dataloader.count();
}

fn main() {
    let d1 = download_dataset(url, "../data", true, true, transform);
    let d2 = download_dataset(url, "../data", false, false, transform);

    let train_dl = DataLoader::new(&d1);
    let test_dl = DataLoader::new(&d2);

    let model = Net::init();
    let optimizer = SGD::init(0.001, 0.8);

    for epoch in 1..EPOCHS+1 {
        train(&model, &train_dl, &optimizer, epoch);
        test(&model, &test_dl, &optimizer, epoch);
    }

    if args.SAVE_MODEL {
        model.save_model("mnist_test.kaffe_pt");        
    }
}

GPU 支持

截至目前，对 GPU 训练的支持不会很快出现。尽管如此..转译确实是一件事。

更多示例请见 这里

文档

完整的 API 文档可以在 这里 找到。

特性

• 非常快速
• 常见的张量操作位于 tensor 模块下
• 优化器
• 支持 f32 和 f64
• ReLU、GeLU、PReLU、Sigmoid
• 基本神经网络特性