Change the utils.download_emo_models (#199)

* Change the utils.download_emo_models

Change utils.download_emo_models(config.mirror, model_name, REPO_ID) to utils.download_emo_models(config.mirror, REPO_ID, model_name)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Thymustw
2023-11-27 16:49:00 +08:00
committed by GitHub
parent 2c528ce130
commit badb12585a

View File

@@ -1,162 +1,163 @@
import argparse import argparse
import os import os
from pathlib import Path from pathlib import Path
import librosa import librosa
import numpy as np import numpy as np
import torch import torch
import torch.nn as nn import torch.nn as nn
from torch.utils.data import DataLoader, Dataset from torch.utils.data import Dataset
from tqdm import tqdm from torch.utils.data import DataLoader, Dataset
from transformers import Wav2Vec2Processor from tqdm import tqdm
from transformers.models.wav2vec2.modeling_wav2vec2 import ( from transformers import Wav2Vec2Processor
Wav2Vec2Model, from transformers.models.wav2vec2.modeling_wav2vec2 import (
Wav2Vec2PreTrainedModel, Wav2Vec2Model,
) Wav2Vec2PreTrainedModel,
)
import utils
from config import config import utils
from config import config
class RegressionHead(nn.Module):
r"""Classification head.""" class RegressionHead(nn.Module):
r"""Classification head."""
def __init__(self, config):
super().__init__() def __init__(self, config):
super().__init__()
self.dense = nn.Linear(config.hidden_size, config.hidden_size)
self.dropout = nn.Dropout(config.final_dropout) self.dense = nn.Linear(config.hidden_size, config.hidden_size)
self.out_proj = nn.Linear(config.hidden_size, config.num_labels) self.dropout = nn.Dropout(config.final_dropout)
self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
def forward(self, features, **kwargs):
x = features def forward(self, features, **kwargs):
x = self.dropout(x) x = features
x = self.dense(x) x = self.dropout(x)
x = torch.tanh(x) x = self.dense(x)
x = self.dropout(x) x = torch.tanh(x)
x = self.out_proj(x) x = self.dropout(x)
x = self.out_proj(x)
return x
return x
class EmotionModel(Wav2Vec2PreTrainedModel):
r"""Speech emotion classifier.""" class EmotionModel(Wav2Vec2PreTrainedModel):
r"""Speech emotion classifier."""
def __init__(self, config):
super().__init__(config) def __init__(self, config):
super().__init__(config)
self.config = config
self.wav2vec2 = Wav2Vec2Model(config) self.config = config
self.classifier = RegressionHead(config) self.wav2vec2 = Wav2Vec2Model(config)
self.init_weights() self.classifier = RegressionHead(config)
self.init_weights()
def forward(
self, def forward(
input_values, self,
): input_values,
outputs = self.wav2vec2(input_values) ):
hidden_states = outputs[0] outputs = self.wav2vec2(input_values)
hidden_states = torch.mean(hidden_states, dim=1) hidden_states = outputs[0]
logits = self.classifier(hidden_states) hidden_states = torch.mean(hidden_states, dim=1)
logits = self.classifier(hidden_states)
return hidden_states, logits
return hidden_states, logits
class AudioDataset(Dataset):
def __init__(self, list_of_wav_files, sr, processor): class AudioDataset(Dataset):
self.list_of_wav_files = list_of_wav_files def __init__(self, list_of_wav_files, sr, processor):
self.processor = processor self.list_of_wav_files = list_of_wav_files
self.sr = sr self.processor = processor
self.sr = sr
def __len__(self):
return len(self.list_of_wav_files) def __len__(self):
return len(self.list_of_wav_files)
def __getitem__(self, idx):
wav_file = self.list_of_wav_files[idx] def __getitem__(self, idx):
audio_data, _ = librosa.load(wav_file, sr=self.sr) wav_file = self.list_of_wav_files[idx]
processed_data = self.processor(audio_data, sampling_rate=self.sr)[ audio_data, _ = librosa.load(wav_file, sr=self.sr)
"input_values" processed_data = self.processor(audio_data, sampling_rate=self.sr)[
][0] "input_values"
return torch.from_numpy(processed_data) ][0]
return torch.from_numpy(processed_data)
def process_func(
x: np.ndarray, def process_func(
sampling_rate: int, x: np.ndarray,
model: EmotionModel, sampling_rate: int,
processor: Wav2Vec2Processor, model: EmotionModel,
device: str, processor: Wav2Vec2Processor,
embeddings: bool = False, device: str,
) -> np.ndarray: embeddings: bool = False,
r"""Predict emotions or extract embeddings from raw audio signal.""" ) -> np.ndarray:
model = model.to(device) r"""Predict emotions or extract embeddings from raw audio signal."""
y = processor(x, sampling_rate=sampling_rate) model = model.to(device)
y = y["input_values"][0] y = processor(x, sampling_rate=sampling_rate)
y = torch.from_numpy(y).unsqueeze(0).to(device) y = y["input_values"][0]
y = torch.from_numpy(y).unsqueeze(0).to(device)
# run through model
with torch.no_grad(): # run through model
y = model(y)[0 if embeddings else 1] with torch.no_grad():
y = model(y)[0 if embeddings else 1]
# convert to numpy
y = y.detach().cpu().numpy() # convert to numpy
y = y.detach().cpu().numpy()
return y
return y
def get_emo(path):
wav, sr = librosa.load(path, 16000) def get_emo(path):
device = config.bert_gen_config.device wav, sr = librosa.load(path, 16000)
return process_func( device = config.bert_gen_config.device
np.expand_dims(wav, 0).astype(np.float64), return process_func(
sr, np.expand_dims(wav, 0).astype(np.float),
model, sr,
processor, model,
device, processor,
embeddings=True, device,
).squeeze(0) embeddings=True,
).squeeze(0)
if __name__ == "__main__":
parser = argparse.ArgumentParser() if __name__ == "__main__":
parser.add_argument( parser = argparse.ArgumentParser()
"-c", "--config", type=str, default=config.bert_gen_config.config_path parser.add_argument(
) "-c", "--config", type=str, default=config.bert_gen_config.config_path
parser.add_argument( )
"--num_processes", type=int, default=config.bert_gen_config.num_processes parser.add_argument(
) "--num_processes", type=int, default=config.bert_gen_config.num_processes
args, _ = parser.parse_known_args() )
config_path = args.config args, _ = parser.parse_known_args()
hps = utils.get_hparams_from_file(config_path) config_path = args.config
hps = utils.get_hparams_from_file(config_path)
device = config.bert_gen_config.device
device = config.bert_gen_config.device
model_name = "./emotional/wav2vec2-large-robust-12-ft-emotion-msp-dim"
REPO_ID = "audeering/wav2vec2-large-robust-12-ft-emotion-msp-dim" model_name = "./emotional/wav2vec2-large-robust-12-ft-emotion-msp-dim"
if not Path(model_name).joinpath("pytorch_model.bin").exists(): REPO_ID = "audeering/wav2vec2-large-robust-12-ft-emotion-msp-dim"
utils.download_emo_models(config.mirror, model_name, REPO_ID) if not Path(model_name).joinpath("pytorch_model.bin").exists():
utils.download_emo_models(config.mirror, REPO_ID, model_name)
processor = Wav2Vec2Processor.from_pretrained(model_name)
model = EmotionModel.from_pretrained(model_name).to(device) processor = Wav2Vec2Processor.from_pretrained(model_name)
model = EmotionModel.from_pretrained(model_name).to(device)
lines = []
with open(hps.data.training_files, encoding="utf-8") as f: lines = []
lines.extend(f.readlines()) with open(hps.data.training_files, encoding="utf-8") as f:
lines.extend(f.readlines())
with open(hps.data.validation_files, encoding="utf-8") as f:
lines.extend(f.readlines()) with open(hps.data.validation_files, encoding="utf-8") as f:
lines.extend(f.readlines())
wavnames = [line.split("|")[0] for line in lines]
dataset = AudioDataset(wavnames, 16000, processor) wavnames = [line.split("|")[0] for line in lines]
data_loader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=16) dataset = AudioDataset(wavnames, 16000, processor)
data_loader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=16)
with torch.no_grad():
for i, data in tqdm(enumerate(data_loader), total=len(data_loader)): with torch.no_grad():
wavname = wavnames[i] for i, data in tqdm(enumerate(data_loader), total=len(data_loader)):
emo_path = wavname.replace(".wav", ".emo.npy") wavname = wavnames[i]
if os.path.exists(emo_path): emo_path = wavname.replace(".wav", ".emo.npy")
continue if os.path.exists(emo_path):
emb = model(data.to(device))[0].detach().cpu().numpy() continue
np.save(emo_path, emb) emb = model(data.to(device))[0].detach().cpu().numpy()
np.save(emo_path, emb)
print("Emo vec 生成完毕!") print("Emo vec 生成完毕!")