burn-train 0.22.0-pre.1

Training crate for the Burn framework
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
122
123
124
125
126
127

use crate::metric::{AccuracyInput, PerplexityInput, TopKAccuracyInput};
use crate::metric::{Adaptor, CerInput, LossInput, WerInput, processor::ItemLazy};
use burn_core::tensor::{Device, Int, Tensor, Transaction};

/// Sequence prediction output adapted for multiple metrics.
///
/// Supported metrics:
/// - Accuracy
/// - TopKAccuracy
/// - Perplexity
/// - Loss
/// - CER
/// - WER
#[derive(new)]
pub struct SequenceOutput {
    /// The loss.
    pub loss: Tensor<1>,

    /// Raw logits. Shape: `[batch_size, seq_len, vocab_size]`
    pub logits: Tensor<3>,

    /// Optional predicted token indices. Shape: `[batch_size, seq_length]`.
    /// If not provided, predictions default to argmax of `logits` along the last dimension.
    pub predictions: Option<Tensor<2, Int>>,

    /// The target token indices. Shape: `[batch_size, seq_length]`
    pub targets: Tensor<2, Int>,
}

impl SequenceOutput {
    fn predicted_tokens(&self) -> Tensor<2, Int> {
        match &self.predictions {
            Some(preds) => preds.clone(),
            None => self.logits.clone().argmax(2).squeeze_dim::<2>(2),
        }
    }

    fn flat_logits(&self) -> Tensor<2> {
        let [batch_size, seq_len, vocab_size] = self.logits.dims();
        self.logits
            .clone()
            .reshape([batch_size * seq_len, vocab_size])
    }

    fn flat_targets(&self) -> Tensor<1, Int> {
        let [batch_size, seq_len] = self.targets.dims();
        self.targets.clone().reshape([batch_size * seq_len])
    }
}

impl ItemLazy for SequenceOutput {
    fn sync(self) -> Self {
        let device: Device = Device::flex();

        match self.predictions {
            Some(preds) => {
                let [logits, loss, targets, predictions] = Transaction::default()
                    .register(self.logits)
                    .register(self.loss)
                    .register(self.targets)
                    .register(preds)
                    .execute()
                    .try_into()
                    .expect("Correct amount of tensor data");

                SequenceOutput {
                    logits: Tensor::from_data(logits, &device),
                    loss: Tensor::from_data(loss, &device),
                    targets: Tensor::from_data(targets, &device),
                    predictions: Some(Tensor::from_data(predictions, &device)),
                }
            }
            None => {
                let [logits, loss, targets] = Transaction::default()
                    .register(self.logits)
                    .register(self.loss)
                    .register(self.targets)
                    .execute()
                    .try_into()
                    .expect("Correct amount of tensor data");

                SequenceOutput {
                    logits: Tensor::from_data(logits, &device),
                    loss: Tensor::from_data(loss, &device),
                    targets: Tensor::from_data(targets, &device),
                    predictions: None,
                }
            }
        }
    }
}

impl Adaptor<LossInput> for SequenceOutput {
    fn adapt(&self) -> LossInput {
        LossInput::new(self.loss.clone())
    }
}

impl Adaptor<CerInput> for SequenceOutput {
    fn adapt(&self) -> CerInput {
        CerInput::new(self.predicted_tokens(), self.targets.clone())
    }
}

impl Adaptor<WerInput> for SequenceOutput {
    fn adapt(&self) -> WerInput {
        WerInput::new(self.predicted_tokens(), self.targets.clone())
    }
}

impl Adaptor<AccuracyInput> for SequenceOutput {
    fn adapt(&self) -> AccuracyInput {
        AccuracyInput::new(self.flat_logits(), self.flat_targets())
    }
}

impl Adaptor<TopKAccuracyInput> for SequenceOutput {
    fn adapt(&self) -> TopKAccuracyInput {
        TopKAccuracyInput::new(self.flat_logits(), self.flat_targets())
    }
}

impl Adaptor<PerplexityInput> for SequenceOutput {
    fn adapt(&self) -> PerplexityInput {
        PerplexityInput::new(self.flat_logits(), self.flat_targets())
    }
}