A Rust-native approach to deep learning.
Currently since there's no mature solution for n-dimensional array computing on gpu in rust, we temporarily use the Tensor
and other CUDA toolkit from tch
, which provides Rust bindings for libtorch
. But we won't use high-level parts of it.
Pytorch version
# create tensor X:
inputs = torch.tensor([[1.0], [3.0], [5.0], [4.0], [8.0], [10.0], [2.0], [6.0]])
# create tensor Y:
labels = torch.tensor([[4.0], [10.0], [16.], [13.0], [25.], [31.], [7.], [19.0]])
# build a model with one linear layer:
model = torch.nn.Linear(1, 1)
# create an SGD optimizer:
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
# train the model on (X, Y):
for epoch in range(5000):
# clear the gradients
optimizer.zero_grad()
# go forward, and ...
loss = torch.nn.functional.mse_loss(model(inputs), labels, reduction='mean')
# ... go backward!
loss.backward()
# update the model parameters:
optimizer.step()
# predict with the trained model:
pred = model(inputs)
print("prediction: ", pred)
Raddar version
// create tensor X:
let inputs = tensor!([[1.0], [3.0], [5.0], [4.0], [8.0], [10.0], [2.0], [6.0]]);
// create tensor Y:
let labels = tensor!([[4.0], [10.0], [16.], [13.0], [25.], [31.], [7.], [19.0]]);
// build a model with one linear layer:
let model = LinearBuilder::default().input_dim(1).output_dim(1).build();
// create an SGD optimizer:
let mut optimizer = opt(model.training_parameters(), GradientDescent::new(0.01));
// train the model on (X, Y):
for epoch in 1..=5000 {
// clear the gradients
model.zero_grad();
// go forward, and ...
let loss = model(&inputs).mse_loss(&labels, Reduction::Mean);
// ... go backward!
loss.backward();
// update the model parameters:
optimizer.step();
}
// predict with the trained model:
let pred = model(&inputs);
println!("prediction: {:?}", pred);
This crate requires CUDA and libtorch
support. You need to:
-
Install CUDA 11.3 from NVIDIA. After the installation, you need to set the
TORCH_CUDA_VERSION
environment variable to11.3
, orlibtorch
cannot find your CUDA. -
Install
libtorch
. You can follow the instructions in tch-rs. -
(Only on Windows) Check if you are using a MSVC-based toolchain. The GNU toolchain can NOT successfully compile
torch-sys
on Windows. You can check the current toolchain by runningrustup toolchain list
If not, run
rustup toolchain install nightly-x86_64-pc-windows-msvc rustup toolchain default nightly-x86_64-pc-windows-msvc
to switch the toolchain.
-
You should now be able to run the project. Try
device_test
intests/tch_test.rs
to see if the all settings are correct. -
Include the following in your
Cargo.toml
:[dependencies] raddar = "*"
And you are ready to go! 🚀
See some details in raddar_array.
Raddar provides automatic differentiation for most tensor operations it supports. This is commonly used to train models using gradient descent. The optimization is performed over variables which are created via a declarative macro tensor!(...)
by defining their initializations in this example below. The model
contains only one linear layer, which has input_dim
and output_dim
of one. Once the model has been generated, a optimizer with optimization algorithm of GradientDescent
and scheduler of StepLR
is created. Then on each step of the training loop:
- The forward pass is applied to the data.
- A loss is computed as the mean square error between the model output and the ground truth.
- Finally an optimization step is performed: gradients are computed and variables of
model
are modified accordingly.
use raddar::nn::{LinearBuilder, Trainable};
use raddar::optim::{
GradientDescent, Optimizer,
};
use raddar::tensor;
use tch::Reduction;
fn gradient_descent() {
let inputs = tensor!([[1.0], [3.0], [5.0], [4.0], [8.0], [10.0], [2.0], [6.0]]);
let labels = tensor!([[4.0], [10.0], [16.], [13.0], [25.], [31.], [7.], [19.0]]);
let model = LinearBuilder::default().input_dim(1).output_dim(1).build();
let mut optimizer = Optimizer::new(
model.training_parameters(),
GradientDescent::new(0.01),
Some(StepLRBuilder::default().build()),
);
for epoch in 1..=5000 {
model.zero_grad();
let loss = model(&inputs).mse_loss(&labels, Reduction::Mean);
loss.backward();
optimizer.step();
println!("epoch: {}, loss: {}", epoch, f64::from(loss));
}
println!("final model:");
println!(
"weight: {}, bias: {}",
f64::from(&*model.module().linear_weight.lock()),
f64::from(&*model.module().linear_bias.as_ref().unwrap().lock())
);
}
The declarative macro seq!(...)
can be used to create neural network architectures. The following code defines a simple model with three hidden layers (one linear layer, one LeakyReLU layer, and one linear layer respectively) and a optimizer with optimization algorithm of RMSProp
and scheduler of StepLR
.
use raddar::nn::{LinearBuilder, LeakyReLU, Mod};
use raddar::optim::{
Optimizer, RMSPropBuilder, StepLRBuilder,
};
use raddar::{seq, tensor};
use tch::{no_grad, Device, Kind, Reduction, Tensor};
fn run() {
let inputs = tensor!([[1.0f32], [3.0], [5.0], [4.0], [8.0], [10.0], [2.0], [6.0]])
.to(tch::Device::Cuda(0));
let labels = tensor!([[4.0f32], [10.0], [16.], [13.0], [25.], [31.], [7.], [19.0]])
.to(tch::Device::Cuda(0));
let model = seq!(
LinearBuilder::default().input_dim(1).output_dim(1).build(),
Mod::new(LeakyReLU::new(0.01)),
)
.to(tch::Device::Cuda(0));
model
.module_mut()
.push(LinearBuilder::default().input_dim(1).output_dim(1).build());
let mut optimizer = Optimizer::new(
model.training_parameters(),
RMSPropBuilder::default().build(),
Some(StepLRBuilder::default().build()),
);
for _epoch in 1..=5000 {
model.zero_grad();
let loss = model(&inputs).mse_loss(&labels, Reduction::Mean);
loss.backward();
optimizer.step();
}
model(&inputs).print();
}
TensorDataset::from_tensors
can be used to create a tensor-dataset. Use into_loader
to create a DataLoader
from the corresponding dataset.
use std::sync::Arc;
use raddar::{
dataset::{DataLoaderConfigBuilder, TensorDataset},
tensor, tensor_vec,
};
use tch::Tensor;
fn create_dataset() {
let inputs = tensor_vec![[1.0], [3.0], [5.0], [4.0], [8.0], [10.0], [2.0], [6.0]];
let labels = tensor_vec![[4.0], [10.0], [16.], [13.0], [25.], [31.], [7.], [19.0]];
let dataset = TensorDataset::from_tensors(inputs, labels);
let mut iter = dataset
.into_loader(
DataLoaderConfigBuilder::default()
.batch_size(3)
.build()
.unwrap(),
)
.cycle();
iter.next();
let (batch, _) = iter.next().unwrap();
assert_eq!(batch.size(), [3, 1]);
let (batch, _) = iter.next().unwrap();
assert_eq!(batch.size(), [2, 1]);
let (batch, _) = iter.next().unwrap();
assert_eq!(batch.size(), [3, 1]);
}
In the following code, the built-in resnet50
model is trained on the CIFAR-10 dataset using RMSProp
optimizer with scheduler algorithm of CosineAnnealingLR
. There are also many other built-in model in Raddar, including VGG, AlexNet and DenseNet.
use std::sync::Arc;
use image::DynamicImage;
use linked_hash_map::LinkedHashMap;
use raddar::dataset::{
image_mappings, DataLoaderConfigBuilder, Dataset, DynImageDataset, LoadFromImageFolder,
TensorDataset, UnsupervisedTensorDataset,
};
use raddar::nn::embedding::{Embedding, OneHot};
use raddar::nn::{
resnet50, LeakyReLU, LinearBuilder, Mod, Trainable
};
use raddar::optim::{
cosine_annealing_lr, Optimizer, rmsprop,
};
use raddar::{seq, tensor};
use tch::{no_grad, Device, Kind, Reduction, Tensor};
fn cifar10() {
let num_classes = 10;
let batch_size = 32;
let device = Device::Cuda(0);
let model = resnet50(num_classes).to(device);
let mut optimizer = opt_with_sched(
model.training_parameters(),
rmsprop(0.03, 0.99),
cosine_annealing_lr(200, 100, 0.01),
);
let classes_vec = vec![
(0, "airplane".to_string()),
(1, "automobile".to_string()),
(2, "bird".to_string()),
(3, "cat".to_string()),
(4, "deer".to_string()),
(5, "dog".to_string()),
(6, "frog".to_string()),
(7, "horse".to_string()),
(8, "ship".to_string()),
(9, "truck".to_string()),
];
let classes_map: LinkedHashMap<_, _> = classes_vec.into_iter().collect();
let mut cifar_dataset = TensorDataset::default();
let mut valid_dataset = TensorDataset::default();
for (id, class) in &classes_map {
let train_path = "dataset/cifar10/train/";
let valid_path = "dataset/cifar10/val/";
let id = id.to_owned();
let train_temp_dataset =
DynImageDataset::from_image_folder(&(train_path.to_owned() + class), ())
.map::<DynImageDataset, _>(image_mappings::resize(224, 224))
.map::<UnsupervisedTensorDataset, _>(image_mappings::to_tensor(
DynamicImage::into_rgb32f,
))
.map::<TensorDataset, _>(move |inputs: Arc<Tensor>| {
let new_inputs = inputs.permute(&[2, 0, 1]);
(Arc::new(new_inputs), Arc::new(Tensor::from(id)))
})
.to(device);
let valid_temp_dataset =
DynImageDataset::from_image_folder(&(valid_path.to_owned() + class), ())
.map::<DynImageDataset, _>(image_mappings::resize(224, 224))
.map::<UnsupervisedTensorDataset, _>(image_mappings::to_tensor(
DynamicImage::into_rgb32f,
))
.map::<TensorDataset, _>(move |inputs: Arc<Tensor>| {
let new_inputs = inputs.permute(&[2, 0, 1]);
(Arc::new(new_inputs), Arc::new(Tensor::from(id)))
})
.to(device);
cifar_dataset = cifar_dataset
.into_iter()
.chain(train_temp_dataset.into_iter())
.collect();
valid_dataset = valid_dataset
.into_iter()
.chain(valid_temp_dataset.into_iter())
.collect();
println!("Class {} is loaded.", *class);
}
let cifar_dataloader = cifar_dataset.into_loader(
DataLoaderConfigBuilder::default()
.batch_size(batch_size)
.shuffle(true)
.build()
.unwrap(),
);
let valid_dataloader = valid_dataset.into_loader(
DataLoaderConfigBuilder::default()
.batch_size(batch_size)
.shuffle(true)
.build()
.unwrap(),
);
println!("DataLoader is prepared.");
let onehot = OneHot::new(num_classes);
for _ in 1..50 {
let epoch_loader = cifar_dataloader.clone();
let valid_loader = valid_dataloader.clone();
for (img, label) in epoch_loader {
model.zero_grad();
let output = model(&img.to_kind(Kind::Float));
let loss = output.mse_loss(
&(onehot)(&label.to_kind(Kind::Int64)).to_kind(Kind::Float),
Reduction::Mean,
);
loss.backward();
optimizer.step();
}
model.eval(true);
let mut acc = Tensor::zeros(&[1], (Kind::Float, device));
let mut now_bnum = 0;
for (img, label) in valid_loader {
no_grad(|| {
let output = model(&img.to_kind(Kind::Float));
let output = output.argmax(1, false);
let bacc = output.eq_tensor(&label).mean(Kind::Float);
acc += bacc;
now_bnum += 1;
});
}
acc = acc / now_bnum;
acc.print();
model.train(true);
}
}