fastnn 0.3.3

A robust deep learning library in Rust with 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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

//! Checkpoints that can resume a run: weights, optimizer state, and the step.
//!
//! ```no_run
//! # use fastnn::prelude::*;
//! # use fastnn::serialize::{save_training, load_training};
//! # let model = Sequential::new().add(Linear::new(4, 2));
//! # let total_steps = 10_000;
//! let mut opt = Adam::new(model.parameters(), 1e-3);
//!
//! // Resume if a checkpoint exists.
//! let mut step = load_training(&model, &mut opt, "run.fdl").unwrap_or(0);
//!
//! while step < total_steps {
//!     // ... train ...
//!     step += 1;
//!     if step % 500 == 0 {
//!         save_training(&model, &opt, step, "run.fdl")?;
//!     }
//! }
//! # Ok::<(), fastnn::Error>(())
//! ```

use std::path::Path;

use crate::error::{Error, Result};
use crate::nn::Module;
use crate::optim::{Optimizer, OptimizerState};
use crate::tensor::Tensor;

use super::checkpoint::{Checkpoint, load_tensors, save_tensors};

/// Name prefixes inside the file. `/` can't clash with the `.` in module paths.
const PARAM: &str = "param/";
const BUFFER: &str = "buffer/";
const OPTIM: &str = "optim/";
const STEP: &str = "meta/step";
const OPTIM_STEPS: &str = "meta/optim_steps";

/// Write model weights, optimizer state, and the step counter to `path`.
pub fn save_training(
    model: &dyn Module,
    optimizer: &dyn Optimizer,
    step: u64,
    path: impl AsRef<Path>,
) -> Result<()> {
    let mut file = Checkpoint::new();

    for (name, param) in model.named_parameters() {
        file.insert(format!("{PARAM}{name}"), param.value().cpu());
    }
    for (name, buffer) in model.named_buffers() {
        file.insert(format!("{BUFFER}{name}"), buffer.value().cpu());
    }

    let state = optimizer.state();
    for (name, tensor) in state.tensors {
        file.insert(format!("{OPTIM}{name}"), tensor);
    }
    file.insert(STEP.into(), Tensor::scalar(step as f32));
    file.insert(OPTIM_STEPS.into(), Tensor::scalar(state.steps as f32));

    save_tensors(&file, path)
}

/// Restore what [`save_training`] wrote and return the step it stopped at.
///
/// A plain model checkpoint is an error, since resuming from it would silently
/// reset the optimizer.
pub fn load_training(
    model: &dyn Module,
    optimizer: &mut dyn Optimizer,
    path: impl AsRef<Path>,
) -> Result<u64> {
    let file = load_tensors(path)?;
    if !file.contains_key(STEP) {
        return Err(Error::Checkpoint(
            "not a training checkpoint; it has no optimizer state, so resuming \
             from it would restart the optimizer. Use load() if that is what you want."
                .into(),
        ));
    }

    for (name, param) in model.named_parameters() {
        let value = expect(&file, &format!("{PARAM}{name}"), &param.shape())?;
        param.set_value(value.to(param.device()));
    }
    for (name, buffer) in model.named_buffers() {
        let device = buffer.value().device();
        let value = expect(&file, &format!("{BUFFER}{name}"), buffer.value().shape())?;
        buffer.set_value(value.to(device));
    }

    let mut state = OptimizerState::new(scalar(&file, OPTIM_STEPS)? as u64);
    for (name, tensor) in &file {
        if let Some(slot) = name.strip_prefix(OPTIM) {
            state.tensors.insert(slot.to_string(), tensor.clone());
        }
    }
    optimizer.load_state(state)?;

    Ok(scalar(&file, STEP)? as u64)
}

fn expect(file: &Checkpoint, name: &str, shape: &[usize]) -> Result<Tensor> {
    let tensor = file
        .get(name)
        .ok_or_else(|| Error::Checkpoint(format!("checkpoint has no '{name}'")))?;
    if tensor.shape() != shape {
        return Err(Error::Checkpoint(format!(
            "'{name}' is {:?} in the file but {shape:?} in the model",
            tensor.shape()
        )));
    }
    Ok(tensor.clone())
}

fn scalar(file: &Checkpoint, name: &str) -> Result<f32> {
    file.get(name)
        .map(|t| t.item())
        .ok_or_else(|| Error::Checkpoint(format!("checkpoint has no '{name}'")))
}