Package: audax
audaxSponsorsAboutInstallationData pipelineWhat's availableAudio feature extractionNetwork architecturesLearnable FrontendsSelf-supervised modelsWhat's coming upOn contributingReferencesSponsorsThis work would not be possible without cloud resources provided by Google'sTPU Research Cloud (TRC) program. I also thank the TRC support team for quickly resolving whatever issues I had: you're awesome!Want to become a sponsor? Feel free to reach out!AboutA home for audio ML inJAX. Has common features, popular learnable frontends, and pretrained supervised and self-supervised models.
As opposed to popular frameworks, the objective is not to become an end-to-end, end-all-be-all DL framework, but instead to act as a starting point for doing things the jax way, through reference implementations and recipes, using thejax/flax/optaxstack.PS:I'm quite new to using Jax and it's functional-at-heart design, so I admit the code can be a bit untidy at places.
Expect changes, restructuring, and like the official Jax repository itself says, sharp edges!InstallationpipinstallaudaxTo install from the latest source use following commandgitclonehttps://github.com/SarthakYadav/audax.gitcdaudax
pipinstall-rrequirements.txt
pipinstall.Data pipelineAll training is done on custom TFRecords. Initially tried usingtensorflow-datasets, but decided against it.tfrecords comprise of examples with audio file stored as an encodedPCM_16flacbuffer, label info and duration, resulting in smallertfrecordfiles and faster I/O as compared to storing audio as a sequence of floats.A step-by-step guide to setup data can be found in therecipes/data_prep, including sample script to convert data into tfrecords.More info could be found inaudax.training_utils.data_v2What's availableAudio feature extractionAt the time of writing,jax.signaldoes not have a native Short-time Fourier Transform (stft) implementation.Instead of trying to emulate thescipy.signalimplementation that has a lot more bells and whistles and is more feature packed,
thestftimplementation inaudax.coreis designed such that it can be build upon to extractspectrogramandmelspectrogramfeatures
as those found intorchaudio, which are quite popular.
The result is a simple implementation ofstft,spectrogramandmelspectrogram, which are compatible with their torchaudio counterparts, as shown in the figure below.Currently,spectrogramandmelspectrogramfeatures are supported. Visitaudax.core.readmefor more info.Apart from features, jax.vmap compatiblemixupandSpecAugment(no TimeStretch as of now unfortunately) implementations are also provided.Network architecturesSeveral prominent neural network architecture reference implementations are provided, with more to come. The current release has:ResNets[1]EfficientNet[2]ConvNeXT[3]Pretrained models can be found in respectiverecipes, and expect more to be added soon.Learnable frontendsTwo popular learnable feature extraction frontends are available inaudax.frontendsLEAF[4] andSincNet[5].
Sample recipes, as well as pretrained models (AudioSetfor now) can be found in therecipes/leaf.Self-supervised modelsCOLA[6] models on AudioSet for various aforementioned architectures can be found inrecipes/cola.A working implementation ofSimCLR[7, 8] can be found inrecipes/simclr, and pretrained models will be added soon (experiments ongoing!).What's coming upPretrainedCOLAmodels and linear probe experiments. (VERY SOON!)Better documentation and walk-throughs.PretrainedSimCLRmodels.Recipes for Speaker Recognition onVoxCelebMoreAudioSetpretrained checkpoints for architectures already added.Reference implementations for more neural architectures, esp. Transformer based networks.On contributingAt the time of writing, I've been the sole person involved in development of this work, and quite frankly, would love to have help!Happy to hear from open source contributore, both newbies and experienced, about their experience and needsAlways open to hearing about possible ways to clean up/better structure code.References[1] He, K., Zhang, X., Ren, S. and Sun, J., 2016. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 770-778).[2] Tan, M. and Le, Q., 2019, May. Efficientnet: Rethinking model scaling for convolutional neural networks. In International conference on machine learning (pp. 6105-6114). PMLR.[3] Liu, Z., Mao, H., Wu, C.Y., Feichtenhofer, C., Darrell, T. and Xie, S., 2022. A ConvNet for the 2020s. arXiv preprint arXiv:2201.03545.[4] Zeghidour, H., Teboul, O., Quitry, F., and Tagliasacchi, M., LEAF: A Learnable Frontend for Audio Classification, In International Conference on Learning Representations, 2021.[5] Ravanelli, M. and Bengio, Y., 2018, December. Speaker recognition from raw waveform with sincnet. In 2018 IEEE Spoken Language Technology Workshop (SLT) (pp. 1021-1028). IEEE.[6] Saeed, A., Grangier, D. and Zeghidour, N., 2021, June. Contrastive learning of general-purpose audio representations. In ICASSP 2021-2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) (pp. 3875-3879). IEEE.[7] Chen, T., Kornblith, S., Norouzi, M. and Hinton, G., 2020, November. A simple framework for contrastive learning of visual representations. In International conference on machine learning (pp. 1597-1607). PMLR.