fastnn 0.3.1

A small deep learning library in Rust with optional CUDA kernels
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91

//! # FastNN
//!
//! A deep learning library in Rust with optional CUDA kernels.
//!
//! ```no_run
//! use fastnn::prelude::*;
//! use fastnn::data::{Mnist, Split};
//!
//! # fn main() -> fastnn::Result<()> {
//! let train = Mnist::load(Split::Train)?;
//! let loader = DataLoader::new(&train, 128).shuffle(true);
//!
//! let model = Sequential::new()
//!     .add(Flatten::new())
//!     .add(Linear::new(784, 128))
//!     .add(ReLU)
//!     .add(Linear::new(128, 10));
//!
//! let mut opt = Adam::new(model.parameters(), 1e-3);
//!
//! for batch in loader.iter() {
//!     let loss = cross_entropy(&model.forward(&batch.inputs), &batch.labels());
//!
//!     opt.zero_grad();
//!     loss.backward();
//!     opt.step();
//! }
//!
//! save(&model, "mnist.fdl")?;
//! # Ok(())
//! # }
//! ```
//!
//! Modules: [`tensor`], [`autograd`], [`nn`], [`optim`], [`data`], [`serialize`], [`cuda`].
//!
//! Tensor maths panics on shape and device mistakes. [`Result`] is for I/O,
//! dataset downloads, and CUDA.
//!
//! Everything starts on the CPU. Move the model with
//! [`Module::to_device`](nn::Module::to_device) and the batches with
//! [`DataLoader::to_device`](data::DataLoader::to_device).
//!
//! ```no_run
//! # use fastnn::prelude::*;
//! # let model = Sequential::new();
//! # let dataset: fastnn::data::TensorDataset = unimplemented!();
//! let device = Device::best();
//! model.to_device(device);
//! let loader = DataLoader::new(&dataset, 64).to_device(device);
//! ```

pub mod autograd;
pub mod cuda;
pub mod data;
pub mod error;
pub mod nn;
pub mod optim;
pub mod rng;
pub mod serialize;
pub mod tensor;

pub use error::{Error, Result};

/// The usual imports for a training script.
///
/// ```
/// use fastnn::prelude::*;
/// ```
pub mod prelude {
    pub use crate::autograd::{detect_anomaly, no_grad};
    pub use crate::data::{Batch, DataLoader, Dataset, TensorDataset};
    pub use crate::nn::{
        bce, bce_with_logits, cross_entropy, huber, kl_divergence, mae, mse, nll,
        AdaptiveAvgPool2d, AvgPool2d, BatchNorm2d, BeamSearch, Buffer, Conv1d, Conv2d,
        ConvTranspose1d, ConvTranspose2d, CrossEntropyLoss, Dropout, Embedding, Flatten, KvCache,
        LayerNorm, LeakyReLU, Linear, MaxPool2d, Module, MultiHeadAttention, Param,
        PositionalEncoding, RMSNorm, ReLU, Reduction, Reshape, Sampler, Sequential, SiLU, Sigmoid,
        Softmax, StackCache, Tanh, Transformer, TransformerBlock, TransformerDecoder,
        TransformerDecoderBlock, TransformerStack, GELU, GRU, LSTM,
    };
    pub use crate::optim::{
        clip_grad_norm, clip_grad_value, Adadelta, Adagrad, Adam, AdamW, Constant, CosineAnnealing,
        Ema, Lion, Lookahead, LrSchedule, OneCycle, Optimizer, OptimizerState, RAdam, RMSprop,
        StepDecay, Warmup, SGD,
    };
    pub use crate::rng::manual_seed;
    pub use crate::serialize::{
        load, load_safetensors, load_training, save, save_safetensors, save_training,
    };
    pub use crate::tensor::{Device, Tensor, Window};
}