codingutf8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert audio spectrogram transformer checkpoints from the original repository url https github comyuangongndast import argparse import json from pathlib import path import torch import torchaudio from datasets import loaddataset from huggingfacehub import hfhubdownload from transformers import astconfig astfeatureextractor astforaudioclassification from transformers utils import logging logging setverbosityinfo logger logging getloggername def getaudiospectrogramtransformerconfigmodelname config astconfig if 1010 in modelname pass elif speechcommands in modelname config maxlength 128 elif 1212 in modelname config timestride 12 config frequencystride 12 elif 1414 in modelname config timestride 14 config frequencystride 14 elif 1616 in modelname config timestride 16 config frequencystride 16 else raise valueerrormodel not supported repoid huggingfacelabelfiles if speechcommands in modelname config numlabels 35 filename speechcommandsv2id2label json else config numlabels 527 filename audiosetid2label json id2label json loadopenhfhubdownloadrepoid filename repotypedataset r id2label intk v for k v in id2label items config id2label id2label config label2id v k for k v in id2label items return config def renamekeyname if module v in name name name replacemodule v audiospectrogramtransformer if clstoken in name name name replaceclstoken embeddings clstoken if disttoken in name name name replacedisttoken embeddings distillationtoken if posembed in name name name replaceposembed embeddings positionembeddings if patchembed proj in name name name replacepatchembed proj embeddings patchembeddings projection transformer blocks if blocks in name name name replaceblocks encoder layer if attn proj in name name name replaceattn proj attention output dense if attn in name name name replaceattn attention self if norm1 in name name name replacenorm1 layernormbefore if norm2 in name name name replacenorm2 layernormafter if mlp fc1 in name name name replacemlp fc1 intermediate dense if mlp fc2 in name name name replacemlp fc2 output dense final layernorm if audiospectrogramtransformer norm in name name name replaceaudiospectrogramtransformer norm audiospectrogramtransformer layernorm classifier head if module mlphead 0 in name name name replacemodule mlphead 0 classifier layernorm if module mlphead 1 in name name name replacemodule mlphead 1 classifier dense return name def convertstatedictorigstatedict config for key in origstatedict copy keys val origstatedict popkey if qkv in key keysplit key split layernum intkeysplit3 dim config hiddensize if weight in key origstatedict faudiospectrogramtransformer encoder layer layernum attention attention query weight val dim origstatedict faudiospectrogramtransformer encoder layer layernum attention attention key weight valdim dim 2 origstatedict faudiospectrogramtransformer encoder layer layernum attention attention value weight valdim else origstatedict faudiospectrogramtransformer encoder layer layernum attention attention query bias val dim origstatedict faudiospectrogramtransformer encoder layer layernum attention attention key bias valdim dim 2 origstatedict faudiospectrogramtransformer encoder layer layernum attention attention value bias valdim else origstatedictrenamekeykey val return origstatedict def removekeysstatedict ignorekeys module v head weight module v head bias module v headdist weight module v headdist bias for k in ignorekeys statedict popk none torch nograd def convertaudiospectrogramtransformercheckpointmodelname pytorchdumpfolderpath pushtohubfalse config getaudiospectrogramtransformerconfigmodelname modelnametourl astfinetunedaudioset10100 4593 https www dropbox comsca0b1v2nlxzyeb4audioset10100 4593 pth dl1 astfinetunedaudioset10100 450 https www dropbox coms1tv0hovue1bxupkaudioset10100 4495 pth dl1 astfinetunedaudioset10100 448 https www dropbox coms6u5sikl4b9wo4u5audioset10100 4483 pth dl1 astfinetunedaudioset10100 448v2 https www dropbox comskt6i0v9fvfm1mbqaudioset10100 4475 pth dl1 astfinetunedaudioset12120 447 https www dropbox comssnfhx3tizr4nuc8audioset12120 4467 pth dl1 astfinetunedaudioset14140 443 https www dropbox comsz18s6pemtnxm4k7audioset14140 4431 pth dl1 astfinetunedaudioset16160 442 https www dropbox comsmdsa4t1xmcimia6audioset16160 4422 pth dl1 astfinetunedspeechcommandsv2 https www dropbox comsq0tbqpwv44pquwyspeechcommands10100 9812 pth dl1 load original statedict checkpointurl modelnametourlmodelname statedict torch hub loadstatedictfromurlcheckpointurl maplocationcpu remove some keys removekeysstatedict rename some keys newstatedict convertstatedictstatedict config load model model astforaudioclassificationconfig model eval model loadstatedictnewstatedict verify outputs on dummy input source https github comyuangongndastblob79e873b8a54d0a3b330dd522584ff2b9926cd581srcrun pyl62 mean 4 2677393 if speechcommands not in modelname else 6 845978 std 4 5689974 if speechcommands not in modelname else 5 5654526 maxlength 1024 if speechcommands not in modelname else 128 featureextractor astfeatureextractormeanmean stdstd maxlengthmaxlength if speechcommands in modelname dataset loaddatasetspeechcommands v0 02 splitvalidation waveform dataset0audioarray else filepath hfhubdownload repoidnielsraudiospectogramtransformercheckpoint filenamesampleaudio flac repotypedataset waveform torchaudio loadfilepath waveform waveform squeeze numpy inputs featureextractorwaveform samplingrate16000 returntensorspt forward pass outputs modelinputs logits outputs logits if modelname astfinetunedaudioset10100 4593 expectedslice torch tensor0 8760 7 0042 8 6602 elif modelname astfinetunedaudioset10100 450 expectedslice torch tensor1 1986 7 0903 8 2718 elif modelname astfinetunedaudioset10100 448 expectedslice torch tensor2 6128 8 0080 9 4344 elif modelname astfinetunedaudioset10100 448v2 expectedslice torch tensor1 5080 7 4534 8 8917 elif modelname astfinetunedaudioset12120 447 expectedslice torch tensor0 5050 6 5833 8 0843 elif modelname astfinetunedaudioset14140 443 expectedslice torch tensor0 3826 7 0336 8 2413 elif modelname astfinetunedaudioset16160 442 expectedslice torch tensor1 2113 6 9101 8 3470 elif modelname astfinetunedspeechcommandsv2 expectedslice torch tensor6 1589 8 0566 8 7984 else raise valueerrorunknown model name if not torch allcloselogits0 3 expectedslice atol1e4 raise valueerrorlogits don t match printlooks ok if pytorchdumpfolderpath is not none pathpytorchdumpfolderpath mkdirexistoktrue printfsaving model modelname to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath printfsaving feature extractor to pytorchdumpfolderpath featureextractor savepretrainedpytorchdumpfolderpath if pushtohub printpushing model and feature extractor to the hub model pushtohubfmitmodelname featureextractor pushtohubfmitmodelname if name main parser argparse argumentparser required parameters parser addargument modelname defaultastfinetunedaudioset10100 4593 typestr helpname of the audio spectrogram transformer model you d like to convert parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the output pytorch model directory parser addargument pushtohub actionstoretrue helpwhether or not to push the converted model to the hub args parser parseargs convertaudiospectrogramtransformercheckpointargs modelname args pytorchdumpfolderpath args pushtohub coding utf 8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert audio spectrogram transformer checkpoints from the original repository url https github com yuangongnd ast transformer blocks final layernorm classifier head copy paste tweak model s weights to our audio spectrogram transformer structure load original state_dict remove some keys rename some keys load model verify outputs on dummy input source https github com yuangongnd ast blob 79e873b8a54d0a3b330dd522584ff2b9926cd581 src run py l62 forward pass required parameters
import argparse
import json
from pathlib import Path

import torch
import torchaudio
from datasets import load_dataset
from huggingface_hub import hf_hub_download

from transformers import ASTConfig, ASTFeatureExtractor, ASTForAudioClassification
from transformers.utils import logging


logging.set_verbosity_info()
logger = logging.get_logger(__name__)


def get_audio_spectrogram_transformer_config(model_name):
    config = ASTConfig()

    if "10-10" in model_name:
        pass
    elif "speech-commands" in model_name:
        config.max_length = 128
    elif "12-12" in model_name:
        config.time_stride = 12
        config.frequency_stride = 12
    elif "14-14" in model_name:
        config.time_stride = 14
        config.frequency_stride = 14
    elif "16-16" in model_name:
        config.time_stride = 16
        config.frequency_stride = 16
    else:
        raise ValueError("Model not supported")

    repo_id = "huggingface/label-files"
    if "speech-commands" in model_name:
        config.num_labels = 35
        filename = "speech-commands-v2-id2label.json"
    else:
        config.num_labels = 527
        filename = "audioset-id2label.json"

    id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
    id2label = {int(k): v for k, v in id2label.items()}
    config.id2label = id2label
    config.label2id = {v: k for k, v in id2label.items()}

    return config


def rename_key(name):
    if "module.v" in name:
        name = name.replace("module.v", "audio_spectrogram_transformer")
    if "cls_token" in name:
        name = name.replace("cls_token", "embeddings.cls_token")
    if "dist_token" in name:
        name = name.replace("dist_token", "embeddings.distillation_token")
    if "pos_embed" in name:
        name = name.replace("pos_embed", "embeddings.position_embeddings")
    if "patch_embed.proj" in name:
        name = name.replace("patch_embed.proj", "embeddings.patch_embeddings.projection")
    
    if "blocks" in name:
        name = name.replace("blocks", "encoder.layer")
    if "attn.proj" in name:
        name = name.replace("attn.proj", "attention.output.dense")
    if "attn" in name:
        name = name.replace("attn", "attention.self")
    if "norm1" in name:
        name = name.replace("norm1", "layernorm_before")
    if "norm2" in name:
        name = name.replace("norm2", "layernorm_after")
    if "mlp.fc1" in name:
        name = name.replace("mlp.fc1", "intermediate.dense")
    if "mlp.fc2" in name:
        name = name.replace("mlp.fc2", "output.dense")
    
    if "audio_spectrogram_transformer.norm" in name:
        name = name.replace("audio_spectrogram_transformer.norm", "audio_spectrogram_transformer.layernorm")
    
    if "module.mlp_head.0" in name:
        name = name.replace("module.mlp_head.0", "classifier.layernorm")
    if "module.mlp_head.1" in name:
        name = name.replace("module.mlp_head.1", "classifier.dense")

    return name


def convert_state_dict(orig_state_dict, config):
    for key in orig_state_dict.copy().keys():
        val = orig_state_dict.pop(key)

        if "qkv" in key:
            key_split = key.split(".")
            layer_num = int(key_split[3])
            dim = config.hidden_size
            if "weight" in key:
                orig_state_dict[
                    f"audio_spectrogram_transformer.encoder.layer.{layer_num}.attention.attention.query.weight"
                ] = val[:dim, :]
                orig_state_dict[
                    f"audio_spectrogram_transformer.encoder.layer.{layer_num}.attention.attention.key.weight"
                ] = val[dim : dim * 2, :]
                orig_state_dict[
                    f"audio_spectrogram_transformer.encoder.layer.{layer_num}.attention.attention.value.weight"
                ] = val[-dim:, :]
            else:
                orig_state_dict[
                    f"audio_spectrogram_transformer.encoder.layer.{layer_num}.attention.attention.query.bias"
                ] = val[:dim]
                orig_state_dict[
                    f"audio_spectrogram_transformer.encoder.layer.{layer_num}.attention.attention.key.bias"
                ] = val[dim : dim * 2]
                orig_state_dict[
                    f"audio_spectrogram_transformer.encoder.layer.{layer_num}.attention.attention.value.bias"
                ] = val[-dim:]
        else:
            orig_state_dict[rename_key(key)] = val

    return orig_state_dict


def remove_keys(state_dict):
    ignore_keys = [
        "module.v.head.weight",
        "module.v.head.bias",
        "module.v.head_dist.weight",
        "module.v.head_dist.bias",
    ]
    for k in ignore_keys:
        state_dict.pop(k, None)


@torch.no_grad()
def convert_audio_spectrogram_transformer_checkpoint(model_name, pytorch_dump_folder_path, push_to_hub=False):
    
    config = get_audio_spectrogram_transformer_config(model_name)

    model_name_to_url = {
        "ast-finetuned-audioset-10-10-0.4593": (
            "https://www.dropbox.com/s/ca0b1v2nlxzyeb4/audioset_10_10_0.4593.pth?dl=1"
        ),
        "ast-finetuned-audioset-10-10-0.450": (
            "https://www.dropbox.com/s/1tv0hovue1bxupk/audioset_10_10_0.4495.pth?dl=1"
        ),
        "ast-finetuned-audioset-10-10-0.448": (
            "https://www.dropbox.com/s/6u5sikl4b9wo4u5/audioset_10_10_0.4483.pth?dl=1"
        ),
        "ast-finetuned-audioset-10-10-0.448-v2": (
            "https://www.dropbox.com/s/kt6i0v9fvfm1mbq/audioset_10_10_0.4475.pth?dl=1"
        ),
        "ast-finetuned-audioset-12-12-0.447": (
            "https://www.dropbox.com/s/snfhx3tizr4nuc8/audioset_12_12_0.4467.pth?dl=1"
        ),
        "ast-finetuned-audioset-14-14-0.443": (
            "https://www.dropbox.com/s/z18s6pemtnxm4k7/audioset_14_14_0.4431.pth?dl=1"
        ),
        "ast-finetuned-audioset-16-16-0.442": (
            "https://www.dropbox.com/s/mdsa4t1xmcimia6/audioset_16_16_0.4422.pth?dl=1"
        ),
        "ast-finetuned-speech-commands-v2": (
            "https://www.dropbox.com/s/q0tbqpwv44pquwy/speechcommands_10_10_0.9812.pth?dl=1"
        ),
    }

    
    checkpoint_url = model_name_to_url[model_name]
    state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")
    
    remove_keys(state_dict)
    
    new_state_dict = convert_state_dict(state_dict, config)

    
    model = ASTForAudioClassification(config)
    model.eval()

    model.load_state_dict(new_state_dict)

    
    
    mean = -4.2677393 if "speech-commands" not in model_name else -6.845978
    std = 4.5689974 if "speech-commands" not in model_name else 5.5654526
    max_length = 1024 if "speech-commands" not in model_name else 128
    feature_extractor = ASTFeatureExtractor(mean=mean, std=std, max_length=max_length)

    if "speech-commands" in model_name:
        dataset = load_dataset("speech_commands", "v0.02", split="validation")
        waveform = dataset[0]["audio"]["array"]
    else:
        filepath = hf_hub_download(
            repo_id="nielsr/audio-spectogram-transformer-checkpoint",
            filename="sample_audio.flac",
            repo_type="dataset",
        )

        waveform, _ = torchaudio.load(filepath)
        waveform = waveform.squeeze().numpy()

    inputs = feature_extractor(waveform, sampling_rate=16000, return_tensors="pt")

    
    outputs = model(**inputs)
    logits = outputs.logits

    if model_name == "ast-finetuned-audioset-10-10-0.4593":
        expected_slice = torch.tensor([-0.8760, -7.0042, -8.6602])
    elif model_name == "ast-finetuned-audioset-10-10-0.450":
        expected_slice = torch.tensor([-1.1986, -7.0903, -8.2718])
    elif model_name == "ast-finetuned-audioset-10-10-0.448":
        expected_slice = torch.tensor([-2.6128, -8.0080, -9.4344])
    elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
        expected_slice = torch.tensor([-1.5080, -7.4534, -8.8917])
    elif model_name == "ast-finetuned-audioset-12-12-0.447":
        expected_slice = torch.tensor([-0.5050, -6.5833, -8.0843])
    elif model_name == "ast-finetuned-audioset-14-14-0.443":
        expected_slice = torch.tensor([-0.3826, -7.0336, -8.2413])
    elif model_name == "ast-finetuned-audioset-16-16-0.442":
        expected_slice = torch.tensor([-1.2113, -6.9101, -8.3470])
    elif model_name == "ast-finetuned-speech-commands-v2":
        expected_slice = torch.tensor([6.1589, -8.0566, -8.7984])
    else:
        raise ValueError("Unknown model name")
    if not torch.allclose(logits[0, :3], expected_slice, atol=1e-4):
        raise ValueError("Logits don't match")
    print("Looks ok!")

    if pytorch_dump_folder_path is not None:
        Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
        print(f"Saving model {model_name} to {pytorch_dump_folder_path}")
        model.save_pretrained(pytorch_dump_folder_path)
        print(f"Saving feature extractor to {pytorch_dump_folder_path}")
        feature_extractor.save_pretrained(pytorch_dump_folder_path)

    if push_to_hub:
        print("Pushing model and feature extractor to the hub...")
        model.push_to_hub(f"MIT/{model_name}")
        feature_extractor.push_to_hub(f"MIT/{model_name}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--model_name",
        default="ast-finetuned-audioset-10-10-0.4593",
        type=str,
        help="Name of the Audio Spectrogram Transformer model you'd like to convert.",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
    )
    parser.add_argument(
        "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
    )

    args = parser.parse_args()
    convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)