autograph
A machine learning library for Rust.
GPGPU kernels implemented with krnl.
- Host and device execution.
- Tensors emulate ndarray
- Host tensors can be borrowed as arrays.
- Tensors, models, and optimizers can be serialized with serde.
- Portable between platforms.
- Save and resume training progress.
- Fully extensible, in Rust.
Neural Networks
let mut model = new?;
model.init_parameter_grads?;
let y = model.forward?;
let loss = y.cross_entropy_loss?;
loss.backward?;
model.update?;
See the Neural Network MNIST example.
Benchmarks
NVIDIA GeForce GTX 1060 with Max-Q Design
LeNet5(training, batch_size = 100)
autograph |
tch |
|
---|---|---|
bf16_host |
482.80 ms (✅ 1.00x) |
75.30 ms (🚀 6.41x faster) |
f32_host |
5.44 ms (✅ 1.00x) |
3.09 ms (✅ 1.76x faster) |
bf16_device |
1.76 ms (✅ 1.00x) |
17.99 ms (❌ 10.20x slower) |
f32_device |
1.75 ms (✅ 1.00x) |
1.20 ms (✅ 1.45x faster) |
LeNet5(inference, batch_size = 1,000)
autograph |
tch |
|
---|---|---|
bf16_host |
1.78 s (✅ 1.00x) |
192.75 ms (🚀 9.25x faster) |
f32_host |
12.23 ms (✅ 1.00x) |
9.57 ms (✅ 1.28x faster) |
bf16_device |
4.62 ms (✅ 1.00x) |
48.72 ms (❌ 10.54x slower) |
f32_device |
4.76 ms (✅ 1.00x) |
1.84 ms (🚀 2.58x faster) |
See the Neural Network benchmark.
License
Dual-licensed to be compatible with the Rust project.
Licensed under the Apache License, Version 2.0 http://www.apache.org/licenses/LICENSE-2.0 or the MIT license http://opensource.org/licenses/MIT, at your option. This file may not be copied, modified, or distributed except according to those terms.
Contribution
Unless you explicitly state otherwise, any contribution intentionally submitted for inclusion in the work by you, as defined in the Apache-2.0 license, shall be dual licensed as above, without any additional terms or conditions.