Dasheng (大声)

Official PyTorch code for Deep Audio-Signal Holistic Embeddings
Scaling up masked audio encoder learning for general audio classification

TL;DR

python3 -m pip install dasheng
python3 -c "from dasheng import dasheng_base; import torch; model = dasheng_base().eval(); features=model(torch.randn(1, 16000))"

This repo provides checkpoints for the Interspeech 2024 paper Scaling up masked audio encoder learning for general audio classification. The goal of this work is to investigate the scalability of masked autoencoders for audio. Prior work did not scale beyond 10,000 hours of audio, while Dasheng used 272,000 hours of training data.

Huggingface 🤗

Please see here for usage instructions.

Models

Dasheng models have been trained on 272k hours of general audio, mainly VGGSound, Audioset, MTG-Jamendo and ACAV100M.

Models with their evaluation results on the HEAR benchmark, averaged across different domains.

K-Nearest Neighbor results

Performance of features without parameterized training.

1. Installation (Recommended for inference)

Install the package.

python3 -m pip install dasheng

1.2 Installation for Training

python3 -m pip install dasheng[train]

2. Usage

# The three models of the paper
from dasheng import dasheng_base, dasheng_06B, dasheng_12B

model = dasheng_base()

Forward some audio data (note should be 16khz)

import torch
model = model.eval()
features = model(torch.randn(1, 16000))
print(features.shape)

3. Training

Install dependencies:

python3 -m pip install dasheng[train]

3.1 Prepare data

We rely on the excellent webdataset library for I/O. Thus one simply needs to pack their data into a bunch of .tar files.

A simple example of such a file would be:

find DIR -type f -name '*flac' |  tar -rvf data.tgz -T -

We also provide a simple script [wavlist_to_tar] that automates this process, which is installed with the package.

wavlist_to_tar your_data.tsv shards/

Creating your_data.tsv is simple:

find data -type f  | awk 'BEGIN{print "filename"} {print}' > your_data.tsv

3.2 Training from source

To train one should first adjust the config in dasheng/train/config/*yaml accordingly, by adding their training data.

python3 dasheng/train/train.py dasheng/train/config/dasheng_base.yaml

MultiGPU support is realized using Accelerate

accelerate launch --mixed_precision='bf16' dasheng/train/train.py dasheng/train/config/dasheng_base.yaml

FAQ

Is there an Audioset-finetuned Dasheng?

Yes, the performance for the base model is 49.7 mAP. One can use it as follows:

from typing import Any, Mapping
import dasheng
import torch

class DashengAudiosetClassifier(torch.nn.Module):

    def __init__(self) -> None:
        super().__init__()
        self.dashengmodel = dasheng.dasheng_base()
        self.classifier = torch.nn.Sequential(torch.nn.LayerNorm(self.dashengmodel.embed_dim), torch.nn.Linear(self.dashengmodel.embed_dim, 527))

    def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True, assign: bool = False):
        self.dashengmodel.load_state_dict(state_dict, strict=False)
        for_classifier_dict = {}
        for k,v in state_dict.items():
            if 'outputlayer' in k:
                for_classifier_dict[k.replace('outputlayer.','')]  = v
        self.classifier.load_state_dict(for_classifier_dict)
        return self

    def forward(self, x):
        x = self.dashengmodel(x).mean(1)
        return self.classifier(x).sigmoid()

mdl = DashengAudiosetClassifier()
check = torch.hub.load_state_dict_from_url('https://zenodo.org/records/13315686/files/dasheng_audioset_mAP497.pt?download=1',map_location='cpu')
mdl.load_state_dict(check)

prediction = mdl(torch.randn(1,16000))

Citation

@inproceedings{dinkel2024dasheng,
  title={Scaling up masked audio encoder learning for general audio classification},
  author={Dinkel, Heinrich and Yan, Zhiyong and Wang, Yongqing and Zhang, Junbo and Wang, Yujun and Wang, Bin},
  booktitle={Interspeech 2024},
  year={2024}
}