dfdx is a cuda accelerated tensor and neural network library, writtten entirely in rust!
Additionally, it can track compile time shapes across tensor operations, ensuring that all your neural networks are checked at compile time.
The following sections provide some high level core concepts & exmaples, and there is more detailed documentation in each of dfdx’s submodules.
See feature_flags for details on feature flags.
At its core a
tensor::Tensor is just a nd-array. Just like
rust arrays there are two parts:
dfdx represents shapes as tuples of dimensions (
where a dimension can either be known at:
You can freely mix and match these dimensions together. Here are some example shapes:
()- unit shape
(usize,)- 1d shape with a runtime known dimension
(usize, Const<5>)- 2d shape with both types of dimensions
(Const<3>, usize, Const<5>)- 3d shape!
Rank3<3, 5, 7>- Equivalent to
(Const<3>, Const<5>, Const<7>)
Here are some comparisons between representing nd arrays in rust vs dfdx:
|rust array||dfdx |
|f32||Tensor<(), f32, …>|
|[u32; 5]||Tensor<Rank1<5>, u32, …>|
|[[u8; 3]; 2]||Tensor<Rank2<2, 3>, u8, …>|
|Vec<[bool; 5]>||Tensor<(usize, Const<5>), bool, …>|
Rank2 shapes used above are actually type aliases for
when all dimensions are compile time:
See tensor for more information.
Devices are used to allocate tensors (and neural networks!). They are akin to std::alloc::GlobalAlloc in rust - they just allocate memory. They are also used to execute tensor ops, which we will get to later on.
There are two options for this currently, with more planned to be added in the future:
Both devices implement Default, you can also create them with a certain seed and ordinal.
Here’s how you might use a device:
let dev: Cpu = Default::default(); let t: Tensor<Rank2<2, 3>, f32, _> = dev.zeros();
See tensor_ops for more information
Once you’ve instantiated tensors with a device, you can start doing operations on them! There are many many operations, here are a few core ones and how they related to things like numpy/pytorch:
|2d Transposed Convolution||tensor_ops::TryConvTrans2D||-|
and much much more!
See nn for more information.
Neural networks are composed of building blocks that you can chain together. In dfdx, sequential neural networks are represents by tuples! For example, the following two networks are identical:
To build a neural network, you of course need a device:
let dev: Cpu = Default::default(); type Model = (Linear<3, 5>, ReLU, Linear<5, 10>); let model = dev.build_module::<Model, f32>();
Note two things:
- We are using nn::DeviceBuildExt to instantiate the model
- We need to pass a dtype (in this case f32) to create the model.
You can then pass tensors into the model with nn::Module::forward():
// tensor with runtime batch dimension of 10 let x: Tensor<(usize, Const<3>), f32, _> = dev.sample_normal_like(&(10, Const)); let y = model.forward(x);
See optim for more information
dfdx supports a number of the standard optimizers:
|AdamW||optim::Adam with optim::WeightDecay::Decoupled|
You can use optimizers to optimize neural networks (or even tensors!). Here’s a simple example of how to do this with nn::ZeroGrads:
type Model = (Linear<3, 5>, ReLU, Linear<5, 10>); let mut model = dev.build_module::<Model, f32>(); // 1. allocate gradients for the model let mut grads = model.alloc_grads(); // 2. create our optimizer let mut opt = Sgd::new(&model, Default::default()); // 3. trace gradients through forward pass let x: Tensor<Rank2<10, 3>, f32, _> = dev.sample_normal(); let y = model.forward_mut(x.traced(grads)); // 4. compute loss & run backpropagation let loss = y.square().mean(); grads = loss.backward(); // 5. apply gradients opt.update(&mut model, &grads);
- Information about the available feature flags.
- Contains subset of all public exports.