burn_train/learner/
sequence.rs

1use crate::metric::{AccuracyInput, PerplexityInput, TopKAccuracyInput};
2use crate::metric::{Adaptor, CerInput, LossInput, WerInput, processor::ItemLazy};
3use burn_core::tensor::{Int, Tensor};
4use burn_std::ExecutionError;
5
6/// Sequence prediction output adapted for multiple metrics.
7///
8/// Supported metrics:
9/// - Accuracy
10/// - TopKAccuracy
11/// - Perplexity
12/// - Loss
13/// - CER
14/// - WER
15#[derive(new)]
16pub struct SequenceOutput {
17    /// The loss.
18    pub loss: Tensor<1>,
19
20    /// Raw logits. Shape: `[batch_size, seq_len, vocab_size]`
21    pub logits: Tensor<3>,
22
23    /// Optional predicted token indices. Shape: `[batch_size, seq_length]`.
24    /// If not provided, predictions default to argmax of `logits` along the last dimension.
25    pub predictions: Option<Tensor<2, Int>>,
26
27    /// The target token indices. Shape: `[batch_size, seq_length]`
28    pub targets: Tensor<2, Int>,
29}
30
31impl SequenceOutput {
32    fn predicted_tokens(&self) -> Tensor<2, Int> {
33        match &self.predictions {
34            Some(preds) => preds.clone(),
35            None => self.logits.clone().argmax(2).squeeze_dim::<2>(2),
36        }
37    }
38
39    fn flat_logits(&self) -> Tensor<2> {
40        let [batch_size, seq_len, vocab_size] = self.logits.dims();
41        self.logits
42            .clone()
43            .reshape([batch_size * seq_len, vocab_size])
44    }
45
46    fn flat_targets(&self) -> Tensor<1, Int> {
47        let [batch_size, seq_len] = self.targets.dims();
48        self.targets.clone().reshape([batch_size * seq_len])
49    }
50}
51
52impl ItemLazy for SequenceOutput {
53    fn sync(self) -> Result<Self, ExecutionError> {
54        // No readback: the metrics compute on the device the tensors live on
55        // and read back only their final scalars, which matters here because
56        // the logits carry the full vocabulary dimension. Flushing dispatches
57        // the producing stream's buffered work so the metric thread doesn't
58        // wait on an idle queue; all tensors in a training item come off the
59        // autodiff backend entirely, so the metric thread neither retains the
60        // tape nor carries its dispatch context.
61        self.loss.device().flush()?;
62
63        Ok(SequenceOutput {
64            logits: self.logits.without_autodiff(),
65            loss: self.loss.without_autodiff(),
66            targets: self.targets.without_autodiff(),
67            predictions: self.predictions.map(|tensor| tensor.without_autodiff()),
68        })
69    }
70}
71
72impl Adaptor<LossInput> for SequenceOutput {
73    fn adapt(&self) -> LossInput {
74        LossInput::new(self.loss.clone())
75    }
76}
77
78impl Adaptor<CerInput> for SequenceOutput {
79    fn adapt(&self) -> CerInput {
80        CerInput::new(self.predicted_tokens(), self.targets.clone())
81    }
82}
83
84impl Adaptor<WerInput> for SequenceOutput {
85    fn adapt(&self) -> WerInput {
86        WerInput::new(self.predicted_tokens(), self.targets.clone())
87    }
88}
89
90impl Adaptor<AccuracyInput> for SequenceOutput {
91    fn adapt(&self) -> AccuracyInput {
92        AccuracyInput::new(self.flat_logits(), self.flat_targets())
93    }
94}
95
96impl Adaptor<TopKAccuracyInput> for SequenceOutput {
97    fn adapt(&self) -> TopKAccuracyInput {
98        TopKAccuracyInput::new(self.flat_logits(), self.flat_targets())
99    }
100}
101
102impl Adaptor<PerplexityInput> for SequenceOutput {
103    fn adapt(&self) -> PerplexityInput {
104        PerplexityInput::new(self.flat_logits(), self.flat_targets())
105    }
106}
burn_train/learner/sequence.rs

burn_train/learner/
sequence.rs
