Remove materialization / Include speech commands dataset [V2.2]

srush / annotated-s4

Implementation of https://srush.github.io/annotated-s4

https://srush.github.io/annotated-s4

MIT License

460 stars 61 forks source link

Remove materialization / Include speech commands dataset [V2.2] #44

Closed srush closed 2 years ago

srush commented 2 years ago

Two more changes for the release

Include the speech commands dataset, a longer dataset of speech outputs.
Fix a bug in JAX that was using up too much memory. Roughly the equivalent of Keops in Pytorch. Now can run on longer sequences with bigger batches (512 model on 6400 time steps with batch 16, will try longer too)

siddk commented 2 years ago

Very cool - LGTM. @srush did you have a date you wanted this to go out? I can do the CIFAR stuff later tomorrow realistically, unclear about W&B integration (or if we need that for the blog).

srush commented 2 years ago

Oh sorry, the wandb feature someone else already added, so that's done.

We could try CIFAR or wait until next release? Not sure if any of these changes will make it better, unless we wanted to try generation.

Currently my speech outputs don't look great, but I guess I've never sampled something this long before :)

siddk commented 2 years ago

I'm fine waiting! Up to you...

Albert/Karan (@albertfgu @krandiash) might have thoughts on how to clean up the audio outputs?

srush commented 2 years ago

Given the dataset is so small, they actually don't look so bad. Think I'm going to declare success.