Type error around "ItemValid = OutputValid" that I just cannot find a way to resolve, if i may beg for help :)

[chrishulbert]

Hi! I've been trying to use Burn to do some ML on time series data, and have been hitting a problem with types that i just cannot resolve.
I was hoping to be able to pick your brains here please :)

Compilation fails on this line: learner.fit(training_dataloader, validation_dataloader)

The error boils down to a type mismatch involving "ItemValid = OutputValid".

I've tried many things to get this to work, such as changing the validation_dataloader to use the Autodiff-wrapped backend, removing explicit types, etc etc but i still get basically the same error.

Could anyone please point me in the right direction?

Here's the full code that i've stripped back to the essentials, and the full error is below that:

#![recursion_limit = "256"] // Burn is too complex for Rust!

use burn::{
    backend::{wgpu::WgpuDevice, Autodiff, Wgpu}, data::{dataloader::{batcher::Batcher, DataLoaderBuilder}, dataset::Dataset}, nn::{loss::CrossEntropyLossConfig, Linear, LinearConfig, Lstm, LstmConfig}, optim::{adaptor::OptimizerAdaptor, Adam, AdamConfig}, prelude::*, tensor::{activation::softmax, backend::AutodiffBackend}, train::{
        ClassificationOutput, LearnerBuilder, TrainOutput, TrainStep, ValidStep
    }
};

// Training details.
const BATCH_SIZE: usize = 64;
const SEQUENCE_LENGTH: usize = 60;
const INPUT_FEATURES: usize = 1;
const OUTPUT_CLASSES: usize = 2;

#[derive(Module, Debug)]
pub struct Model<B: Backend> {
    lstm: Lstm<B>,
    fc: Linear<B>, // Fully connected outputs.
}

#[derive(Config, Debug)]
pub struct ModelConfig {
    #[config(default = 1)] // Must match FEATURES.
    pub inputs: usize,

    #[config(default = 32)]
    pub hidden_size: usize, // LSTM size.

    #[config(default = 2)] // Must match OUTPUT_CLASSES.
    output_classes: usize, // Long, Short, NoTrade.
}

impl ModelConfig {
    /// Returns the initialized model.
    pub fn init<B: Backend>(&self, device: &B::Device) -> Model<B> {
        Model {
            lstm: LstmConfig::new(self.inputs, self.hidden_size, true).init(device),
            fc: LinearConfig::new(self.hidden_size, self.output_classes).init(device),
        }
    }
}

impl<B: Backend> Model<B> {
    // Input shape: [batch_size, sequence_length, num_features]
    // Output: [batch_size, classes] - one-hot logits, not indices yet.
    pub fn forward(&self, input: Tensor<B, 3>) -> Tensor<B, 2> {
        let sequence_length = input.dims()[1];
        let (x, _lstm_state) = self.lstm.forward(input, None); // [batch, seq, hidden]
        let x = x.slice(s![.., (sequence_length-1)..sequence_length, ..]); // Just get the last time: [batch, 1, hidden]
        let x = x.squeeze::<2>(1); // Convert it to 2d. Just get the hiddens for that last time. [batch, hidden]
        let x = self.fc.forward(x); // [batch, num_classes]
        let x = softmax(x, 1); // Dimension index 1 = 'num classes'. [batch, num_classes]
        x
    }
}

fn main() {
    type MyBackend = Wgpu;
    type MyAutodiffBackend = Autodiff<MyBackend>;
    let config = ModelConfig::new();
    let device = WgpuDevice::default();
    let model = config.init::<MyAutodiffBackend>(&device);

    type MyModule = Model<MyAutodiffBackend>;
    type MyOptimizer = OptimizerAdaptor<Adam, MyModule, MyAutodiffBackend>;
    type MyOutput = ClassificationOutput<MyAutodiffBackend>;
    let learner_builder = LearnerBuilder::<
        MyAutodiffBackend,
        MyOutput,
        MyOutput,
        MyModule,
        MyOptimizer,
        f64>::new(".");
    let learner = learner_builder.build(model, AdamConfig::new().init(), 1.0e-3);

    let training_dataset = MyDataset {};
    let validation_dataset = MyDataset {};
    let training_batcher = MyBatcher {};
    let validation_batcher = MyBatcher {};

    let training_dataloader_builder: DataLoaderBuilder<MyAutodiffBackend, MyItem, MyBatch<MyAutodiffBackend>> =
        DataLoaderBuilder::new(training_batcher);
    let training_dataloader = training_dataloader_builder.build(training_dataset);

    let validation_dataloader_builder: DataLoaderBuilder<MyBackend, MyItem, MyBatch<MyBackend>> =
        DataLoaderBuilder::new(validation_batcher);
    let validation_dataloader = validation_dataloader_builder.build(validation_dataset);

    learner.fit(training_dataloader, validation_dataloader);
}

#[derive(Debug)]
struct MyDataset {}
impl Dataset<MyItem> for MyDataset {
    fn get(&self, _index: usize) -> Option<MyItem> { None }
    fn len(&self) -> usize { 0 }
}

#[derive(Debug, Clone)]
struct MyItem {}

#[derive(Clone, Debug)]
struct MyBatch<B: Backend> {
    inputs: Tensor<B, 3>,  // [batch_size, seq_len, features]
    targets: Tensor<B, 1, Int>, // [batch_size] with indexed classes.
}

struct MyBatcher {}
impl<B: Backend> Batcher<B, MyItem, MyBatch<B>> for MyBatcher {
    fn batch(&self, _items: Vec<MyItem>, device: &B::Device) -> MyBatch<B> { // Returns a single batch.
        let input_samples: Vec<f32> = Vec::new();
        let target_samples: Vec<i32> = Vec::new();
        let inputs_d = TensorData::new(input_samples, [0, SEQUENCE_LENGTH, INPUT_FEATURES]);
        let inputs = Tensor::<B,3>::from_data(inputs_d, device);
        let targets_d = TensorData::new(target_samples, [0]);
        let targets = Tensor::<B,1, Int>::from_data(targets_d, device);
        MyBatch { inputs, targets }
    }
}

impl<B: Backend> Model<B> {
    pub fn forward_classification(
        &self,
        inputs: Tensor<B, 3>, // [batch_size, sequence_length, input_features]
        targets: Tensor<B, 1, Int>, // 1D tensor of class indices.
    ) -> ClassificationOutput<B> {
        let logits = self.forward(inputs);
        let loss = CrossEntropyLossConfig::new()
            .init(&logits.device())
            .forward(logits.clone(), targets.clone());
        ClassificationOutput::new(loss, logits, targets)
    }
}

impl<B: AutodiffBackend> TrainStep<MyBatch<B>, ClassificationOutput<B>> for Model<B> {
    fn step(&self, batch: MyBatch<B>) -> TrainOutput<ClassificationOutput<B>> {
        let item = self.forward_classification(batch.inputs, batch.targets);
        TrainOutput::new(self, item.loss.backward(), item)
    }
}

impl<B: Backend> ValidStep<MyBatch<B>, ClassificationOutput<B>> for Model<B> {
    fn step(&self, batch: MyBatch<B>) -> ClassificationOutput<B> {
        self.forward_classification(batch.inputs, batch.targets)
    }
}

Full error:

    Fresh burn-wgpu v0.18.0
    Fresh burn-train v0.18.0
    Fresh burn v0.18.0
Compiling my_burn v0.1.0
    Running `/Users/chris/.rustup/toolchains/stable-aarch64-apple-darwin/bin/rustc --crate-name my_burn --edition=2024 src/main.rs --error-format=json --json=diagnostic-rendered-ansi,artifacts,future-incompat --diagnostic-width=154 --crate-type bin --emit=dep-info,link -C embed-bitcode=no -C debuginfo=2 -C split-debuginfo=unpacked --check-cfg 'cfg(docsrs,test)' --check-cfg 'cfg(feature, values())' -C metadata=9a720066cb167aaa -C extra-filename=-ee4391586fe28682 --out-dir /Users/chris/dev/x/burn/target/debug/deps -C incremental=/Users/chris/dev/x/burn/target/debug/incremental -L dependency=/Users/chris/dev/x/burn/target/debug/deps --extern burn=/Users/chris/dev/x/burn/target/debug/deps/libburn-dfa3ec4ae7511944.rlib --extern rand=/Users/chris/dev/x/burn/target/debug/deps/librand-16956bec85c3a563.rlib -L native=/Users/chris/dev/x/burn/target/debug/build/libsqlite3-sys-6dfa5876f8252cb7/out`
error[E0271]: type mismatch resolving `<AsyncProcessor<...> as EventProcessor>::ItemValid == ClassificationOutput<...>`
--> src/main.rs:89:13
    |
89  |     learner.fit(training_dataloader, validation_dataloader);
    |             ^^^ expected `ClassificationOutput<Fusion<...>>`, found `ClassificationOutput<Autodiff<...>>`
    |
    = note: expected struct `ClassificationOutput<burn_fusion::backend::Fusion<_>>`
            found struct `ClassificationOutput<Autodiff<burn_fusion::backend::Fusion<_>>>`
note: required by a bound in `burn_train::learner::train_val::<impl Learner<LC>>::fit`
--> /Users/chris/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-train-0.18.0/src/learner/train_val.rs:124:69
    |
112 |     pub fn fit<InputTrain, InputValid, OutputTrain, OutputValid>(
    |            --- required by a bound in this associated function
...
124 |         LC::EventProcessor: EventProcessor<ItemTrain = OutputTrain, ItemValid = OutputValid>,
    |                                                                     ^^^^^^^^^^^^^^^^^^^^^^^ required by this bound in `burn_train::learner::train_val::<impl Learner<LC>>::fit`
    = note: the full name for the type has been written to '/Users/chris/dev/x/burn/target/debug/deps/my_burn-ee4391586fe28682.long-type-17011300534260092292.txt'
    = note: consider using `--verbose` to print the full type name to the console

For more information about this error, try rustc --explain E0271.
error: could not compile my_burn (bin "my_burn") due to 1 previous error

Thanks!

[chrishulbert]

Also, i tried following the discord link at the top of burn.dev however it says the invite is invalid/may be expired. I'd love to join if possible! Thanks :)

[laggui]

Are you still having troubles? The invite has no expiration iirc, last member joined 10hrs ago.

[chrishulbert]

Ah, i got it to work by joining the server in the discord desktop app, as opposed to clicking the link on the web. Sounds like discord's being funny, not your invite expired. Thanks!

[chrishulbert]

I was able to resolve this by starting from the 'guide' sample code and copy in my file one function at a time and get it to work that way. However I never really understood the original error, i'm curious if anyone has advice. Maintainers may feel free to close this however as the immediate need is over. Thanks!

[laggui]

The compiler is usually able to infer most types based on the types you provide, but in this case you actually specified each generic explicitly and the valid output was incorrect:

    type MyModule = Model<MyAutodiffBackend>;
    type MyOptimizer = OptimizerAdaptor<Adam, MyModule, MyAutodiffBackend>;
    type MyOutput = ClassificationOutput<MyAutodiffBackend>;
+   type MyOutputVal = ClassificationOutput<MyBackend>;
    let learner_builder = LearnerBuilder::<
        MyAutodiffBackend,
        MyOutput,
-       MyOutput,
+       MyOutputVal,
        MyModule,
        MyOptimizer,
        f64>::new(".");
    let learner = learner_builder.build(model, AdamConfig::new().init(), 1.0e-3);

The validation output should not be on the autodiff backend 🙂

[chrishulbert]

Thankyou so so much! I made the types explicit in my efforts to make it compile when implementing it gradually, because i hadn't entered the lower-down code that would have completed the inference. Good to know!