Repository files navigation

Ergonomics & safety focused deep learning in Rust.

Still in pre-alpha state. The next few releases are planned to be breaking releases.

Features at a glance:

1. 🔥 GPU accelerated tensor library with shapes up to 6d!
2. Shapes with both compile and runtime sized dimensions. (e.g.Tensor<(usize, Const<10>)> andTensor<Rank2<5, 10>>)
3. A large library of tensor operations (includingmatmul,conv2d, and much more).
   1. All tensor operations shape and type checked at compile time!!
4. Ergonomic neural network building blocks (likeLinear,Conv2D, andTransformer).
5. Standard deep learning optimizers such asSgd,Adam,AdamW,RMSprop, and more.

dfdx is oncrates.io! Use by adding this to yourCargo.toml:

dfdx ="0.13.0"

See the documentation atdocs.rs/dfdx.

[1]https://en.wikipedia.org/wiki/Automatic_differentiation#Reverse_accumulation

1. Ergonomics the whole way down (both frontend interface & internals).
2. Check as much at compile time as possible (i.e. don't compile if something is not correct).
3. Maximize performance.
4. Minimize unsafe code[1]
5. Minimize Rc<RefCell> used in internal code[2]

[1] Currently the only unsafe calls are for matrix multiplication.

[2] The only things that useArc are tensors to store their data.Arc is used instead ofBox to reduceallocations when tensors are cloned.

Enable thecuda feature to start using theCuda device! Requires the installation of nvidia's cuda toolkit. Seefeature flags docs for more info.

Checkexamples/ for more details.

1. 👌 Simple Neural Networks API, completely shape checked at compile time.

typeMlp =((Linear<10,32>,ReLU),(Linear<32,32>,ReLU),(Linear<32,2>,Tanh),);fnmain(){let dev:Cuda =Default::default();// or `Cpu`let mlp = dev.build_module::<Mlp,f32>();let x:Tensor<Rank1<10>,f32,Cpu> = dev.zeros();let y:Tensor<Rank1<2>,f32,Cpu> = mlp.forward(x);    mlp.save("checkpoint.npz")?;}

2. 📈 Ergonomic Optimizer API

typeModel = ...letmut model = dev.build_module::<Model,f32>();letmut grads = model.alloc_grads();letmut sgd =Sgd::new(&model,SgdConfig{lr:1e-2,momentum:Some(Momentum::Nesterov(0.9))});let loss = ...grads = loss.backward();sgd.update(&mut model,&grads);

3. 💡 Const tensors can be converted to and from normal rust arrays

let t0:Tensor<Rank0,f32,_> = dev.tensor(0.0);assert_eq!(t0.array(),&0.0);let t1/*: Tensor<Rank1<3>, f32, _>*/ = dev.tensor([1.0,2.0,3.0]);assert_eq!(t1.array(),[1.0,2.0,3.0]);let t2:Tensor<Rank2<2,3>,f32,_> = dev.sample_normal();assert_ne!(t2.array(),[[0.0;3];2]);

pubtraitModule<Input>{typeOutput;fnforward(&self,input:Input) ->Self::Output;}

From this flexible trait we get:

1. Single & batched inputs (just have multiple impls!)
2. Multiple inputs/outputs (multi-headed modules, or rnns)
3. Behavior different when tape is present or not (not the .train()/.eval() behavior present in other libraries!).

Since we can implement traits for tuples, which isnot possible in other languages AFAIK, they provide a very nice frontendfor sequentially executing modules.

// no idea why you would do this, but you could!typeModel =(ReLU,Sigmoid,Tanh);let model = dev.build_module::<Model,f32>();

typeModel =(Linear<10,5>,Tanh)let model = dev.build_module::<Model,f32>();

How implementing Module for a 2-tuple looks:

impl<Input,A,B>Module<Input>for(A,B)whereInput:Tensor,A:Module<Input>,// A is a module that takes InputB:Module<A::Output>,// B is a module that takes A's Output{typeOutput =B::Output;// the output of this is B's Outputfnforward(&self,x:Input) ->Self::Output{let x =self.0.forward(x);let x =self.1.forward(x);        x}}

Modules implemented for Tuples up to 6 elements, butyou can arbitrarily nest them!

Other implementations may store a reference to the gradient tape directly on tensors, which requires mutating tensors or using Rc/Refcells all over the place.

We've figured out an elegant way to avoid this, reducing references and dynamic borrow checks to 0!

Since all operations result in exactly 1 child, we can always move the gradient tape to the child of the last operation. Additionally, no model parameters (all tensors) will ever own the gradient tape because they will never be the result of any operation. This means we know exactly which tensor owns the gradient tape, and the tensors that have it will always be intermediate results that don't need to be maintained across gradient computation.

All of this together gives users unprecedented control/precision over what tensors are recorded on the gradient tape!

One advanced use case requires that tensors be re-used multiple times in a computation graph.This can be handled by cloning the tensor, and manually moving the gradient tape around.

tl;dr: If you forget to include a call totrace() ortraced(), the program won't compile!

-let pred = module.forward(x);+let pred = module.forward(x.traced(grads));let loss = (y - pred).square().mean();let gradients = loss.backward();

Since we know exactly what tensors own the gradient tape, we can require the tensor passed into.backward() to own the gradient tape!And further, we can require it be moved into.backward(), so it can destruct the tape and construct the gradients!

All of this can be checked at compile time 🎉

All functions & operations are tested against behavior shown by similar code in pytorch.

Dual-licensed to be compatible with the Rust project.

Licensed under the Apache License, Version 2.0http://www.apache.org/licenses/LICENSE-2.0 or the MIT licensehttp://opensource.org/licenses/MIT, at your option. This file may not be copied, modified, or distributed except according to those terms.