burn-core 0.23.0-pre.1

Flexible and Comprehensive Deep Learning Framework in Rust
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

use alloc::vec::Vec;

use super::DistributedLayer;
use crate::module::Module;

/// A model that runs as an input layer, then hidden layers in order, then an output layer, so its
/// layers can be split across devices by a [`DistributedLayeredModel`](super::DistributedLayeredModel).
///
/// It is usually not the struct the model trained with, since a layer split across devices may be
/// represented differently. Weights trained on another struct load into this one by remapping
/// their keys. Each layer is built on the device a [`LayerPlacement`](super::LayerPlacement) gives
/// it, so the model never exists whole on one device.
///
/// ```rust,ignore
/// impl Model {
///     fn new(config: &ModelConfig, placement: &LayerPlacement) -> Self {
///         Self {
///             embedding: Embedding::new(config, &placement.input),
///             blocks: placement
///                 .hidden
///                 .iter()
///                 .map(|device| Block::new(config, device))
///                 .collect(),
///             head: Head::new(config, &placement.output),
///         }
///     }
/// }
///
/// impl LayerParallelism for Model {
///     type InputLayer = Embedding;
///     type HiddenLayer = Block;
///     type OutputLayer = Head;
///
///     fn layer_input(&self) -> &Embedding {
///         &self.embedding
///     }
///
///     fn layers_hidden(&self) -> impl Iterator<Item = &Block> {
///         self.blocks.iter()
///     }
///
///     fn layer_output(&self) -> &Head {
///         &self.head
///     }
///
///     fn into_layers(self) -> LayerSequence<Self> {
///         LayerSequence {
///             input: self.embedding,
///             hidden: self.blocks,
///             output: self.head,
///         }
///     }
/// }
///
/// let placement = LayerPlacement::even(&devices, config.num_blocks);
/// let model = Model::new(&config, &placement);
/// let model = DistributedLayeredModel::new(model, &placement);
/// let predictions = model.forward(features);
/// ```
pub trait LayerParallelism: Module {
    /// Runs first, on what the model takes; what it returns is what the hidden layers pass along.
    type InputLayer: DistributedLayer<Input: Module, Output: Module>;
    /// Each hidden layer.
    type HiddenLayer: DistributedLayer<Input = HiddenLayerSignal<Self>, Output = HiddenLayerSignal<Self>>;
    /// Runs last; what it returns is what the model returns.
    type OutputLayer: DistributedLayer<Input = HiddenLayerSignal<Self>>;

    /// The input layer.
    fn layer_input(&self) -> &Self::InputLayer;

    /// The hidden layers in forward order, the same order on every call.
    fn layers_hidden(&self) -> impl Iterator<Item = &Self::HiddenLayer>;

    /// The output layer.
    fn layer_output(&self) -> &Self::OutputLayer;

    /// The model taken apart into its layers, the hidden ones in the order
    /// [`layers_hidden`](Self::layers_hidden) gives them.
    fn into_layers(self) -> LayerSequence<Self>;
}

/// What the hidden layers of a [`LayerParallelism`] model take and return: whatever its input
/// layer produces.
pub type HiddenLayerSignal<M> = <<M as LayerParallelism>::InputLayer as DistributedLayer>::Output;

/// The layers of a [`LayerParallelism`] model, owned.
#[derive(Debug)]
pub struct LayerSequence<M: LayerParallelism> {
    /// The input layer.
    pub input: M::InputLayer,
    /// The hidden layers, in forward order.
    pub hidden: Vec<M::HiddenLayer>,
    /// The output layer.
    pub output: M::OutputLayer,
}