Merge pull request #526 from snakers4/adamnsandle

код для тюнинга

README.md

@@ -120,7 +120,7 @@ Please see our [wiki](https://github.com/snakers4/silero-models/wiki) for releva
 @misc{Silero VAD,
   author = {Silero Team},
   title = {Silero VAD: pre-trained enterprise-grade Voice Activity Detector (VAD), Number Detector and Language Classifier},
-  year = {2021},
+  year = {2024},
   publisher = {GitHub},
   journal = {GitHub repository},
   howpublished = {\url{https://github.com/snakers4/silero-vad}},

tuning/README.md

@@ -0,0 +1,74 @@
+# Тюнинг Silero-VAD модели
+
+> Код тюнинга создан при поддержке Фонда содействия инновациям в рамках федерального проекта «Искусственный
+интеллект» национальной программы «Цифровая экономика Российской Федерации».
+
+Тюнинг используется для улучшения качества детекции речи Silero-VAD модели на кастомных данных.
+
+## Зависимости
+Следующие зависимости используются при тюнинге VAD модели:
+- `torchaudio>=0.12.0`
+- `omegaconf>=2.3.0`
+- `sklearn>=1.2.0`
+- `torch>=1.12.0`
+- `pandas>=2.2.2`
+- `tqdm`
+
+## Подготовка данных
+
+Датафреймы для тюнинга должны быть подготовлены и сохранены в формате `.feather`. Следующие колонки в `.feather` файлах тренировки и валидации являются обязательными:
+- **audio_path** - абсолютный путь до аудиофайла в дисковой системе. Аудиофайлы должны представлять собой `PCM` данные, предпочтительно в форматах `.wav` или `.opus` (иные популярные форматы аудио тоже поддерживаются). Для ускорения темпа дообучения рекомендуется предварительно выполнить ресемплинг аудиофайлов (изменить частоту дискретизации) до 16000 Гц;
+- **speech_ts** - разметка для соответствующего аудиофайла. Список, состоящий из словарей формата `{'start': START_SEC, 'end': 'END_SEC'}`, где `START_SEC` и `END_SEC` - время начало и конца речевого отрезка в секундах соответственно. Для качественного дообучения рекомендуется использовать разметку с точностью до 30 миллисекунд.
+
+Чем больше данных используется на этапе дообучения, тем эффективнее показывает себя адаптированная модель на целевом домене. Длина аудио не ограничена, т.к. каждое аудио будет обрезано до `max_train_length_sec` секунд перед подачей в нейросеть. Длинные аудио лучше предварительно порезать на кусочки длины `max_train_length_sec`.
+
+Пример `.feather` датафрейма можно посмотреть в файле `example_dataframe.feather`
+
+## Файл конфигурации `config.yml`
+
+Файл конфигурации `config.yml` содержит пути до обучающей и валидационной выборки, а также параметры дообучения:
+- `train_dataset_path` - абсолютный путь до тренировочного датафрейма в формате `.feather`. Должен содержать колонки `audio_path` и `speech_ts`, описанные в пункте "Подготовка данных". Пример устройства датафрейма можно посмотреть в `example_dataframe.feather`;
+- `val_dataset_path` - абсолютный путь до валидационного датафрейма в формате `.feather`. Должен содержать колонки `audio_path` и `speech_ts`, описанные в пункте "Подготовка данных". Пример устройства датафрейма можно посмотреть в `example_dataframe.feather`;
+- `jit_model_path` - абсолютный путь до Silero-VAD модели в формате `.jit`. Если оставить это поле пустым, то модель будет загружена из репозитория в зависимости от значения поля `use_torchhub`
+- `use_torchhub` - Если `True`, то модель для дообучения будет загружена с помощью torch.hub. Если `False`, то модель для дообучения будет загружена с помощью библиотеки silero-vad (необходимо заранее установить командой `pip install silero-vad`);
+- `tune_8k` - данный параметр отвечает, какую голову Silero-VAD дообучать. Если `True`, дообучаться будет голова с 8000 Гц частотой дискретизации, иначе с 16000 Гц;
+- `model_save_path` - путь сохранения добученной модели;
+- `noise_loss` - коэффициент лосса, применяемый для неречевых окон аудио;
+- `max_train_length_sec` - максимальная длина аудио в секундах на этапе дообучения. Более длительные аудио будут обрезаны до этого показателя;
+- `aug_prob` - вероятность применения аугментаций к аудиофайлу на этапе дообучения;
+- `learning_rate` - темп дообучения;
+- `batch_size` - размер батча при дообучении и валидации;
+- `num_workers` - количество потоков, используемых для загрузки данных;
+- `num_epochs` - количество эпох дообучения. За одну эпоху прогоняются все тренировочные данные;
+- `device` - `cpu` или `cuda`.
+
+## Дообучение
+
+Дообучение запускается командой 
+
+`python tune.py`
+
+Длится в течение `num_epochs`, лучший чекпоинт по показателю ROC-AUC на валидационной выборке будет сохранен в `model_save_path` в формате jit.
+
+## Поиск пороговых значений
+
+Порог на вход и порог на выход можно подобрать, используя команду 
+
+`python search_thresholds`
+
+Данный скрипт использует файл конфигурации, описанный выше. Указанная в конфигурации модель будет использована для поиска оптимальных порогов на валидационном датасете.
+
+## Цитирование
+
+```
+@misc{Silero VAD,
+  author = {Silero Team},
+  title = {Silero VAD: pre-trained enterprise-grade Voice Activity Detector (VAD), Number Detector and Language Classifier},
+  year = {2024},
+  publisher = {GitHub},
+  journal = {GitHub repository},
+  howpublished = {\url{https://github.com/snakers4/silero-vad}},
+  commit = {insert_some_commit_here},
+  email = {hello@silero.ai}
+}
+```

tuning/config.yml

@@ -0,0 +1,17 @@
+jit_model_path: ''  # путь до Silero-VAD модели в формате jit, эта модель будет использована для дообучения. Если оставить поле пустым, то модель будет загружена автоматически
+use_torchhub: True  # jit модель будет загружена через torchhub, если True, или через pip, если False
+
+tune_8k: False  # дообучает 16к голову, если False, и 8к голову, если True
+train_dataset_path: 'train_dataset_path.feather'  # путь до датасета в формате feather для дообучения, подробности в README
+val_dataset_path: 'val_dataset_path.feather'  # путь до датасета в формате feather для валидации, подробности в README
+model_save_path: 'model_save_path.jit'  # путь сохранения дообученной модели
+
+noise_loss: 0.5  # коэффициент, применяемый к лоссу на неречевых окнах
+max_train_length_sec: 8  # во время тюнинга аудио длиннее будут обрезаны до данного значения
+aug_prob: 0.4  # вероятность применения аугментаций к аудио в процессе дообучения
+
+learning_rate: 5e-4  # темп дообучения модели
+batch_size: 128  # размер батча при дообучении и валидации
+num_workers: 4  # количество потоков, используемых для даталоадеров
+num_epochs: 20  # количество эпох дообучения, 1 эпоха = полный прогон тренировочных данных
+device: 'cuda'  # cpu или cuda, на чем будет производится дообучение

tuning/search_thresholds.py

@@ -0,0 +1,36 @@
+from utils import init_jit_model, predict, calculate_best_thresholds, SileroVadDataset, SileroVadPadder
+from omegaconf import OmegaConf
+import torch
+torch.set_num_threads(1)
+
+if __name__ == '__main__':
+    config = OmegaConf.load('config.yml')
+
+    loader = torch.utils.data.DataLoader(SileroVadDataset(config, mode='val'),
+                                         batch_size=config.batch_size,
+                                         collate_fn=SileroVadPadder,
+                                         num_workers=config.num_workers)
+
+    if config.jit_model_path:
+        print(f'Loading model from the local folder: {config.jit_model_path}')
+        model = init_jit_model(config.jit_model_path, device=config.device)
+    else:
+        if config.use_torchhub:
+            print('Loading model using torch.hub')
+            model, _ = torch.hub.load(repo_or_dir='snakers4/silero-vad',
+                                      model='silero_vad',
+                                      onnx=False,
+                                      force_reload=True)
+        else:
+            print('Loading model using silero-vad library')
+            from silero_vad import load_silero_vad
+            model = load_silero_vad(onnx=False)
+
+    print('Model loaded')
+    model.to(config.device)
+
+    print('Making predicts...')
+    all_predicts, all_gts = predict(model, loader, config.device, sr=8000 if config.tune_8k else 16000)
+    print('Calculating thresholds...')
+    best_ths_enter, best_ths_exit, best_acc = calculate_best_thresholds(all_predicts, all_gts)
+    print(f'Best threshold: {best_ths_enter}\nBest exit threshold: {best_ths_exit}\nBest accuracy: {best_acc}')

tuning/tune.py

@@ -0,0 +1,65 @@
+from utils import SileroVadDataset, SileroVadPadder, VADDecoderRNNJIT, train, validate, init_jit_model
+from omegaconf import OmegaConf
+import torch.nn as nn
+import torch
+
+
+if __name__ == '__main__':
+    config = OmegaConf.load('config.yml')
+
+    train_dataset = SileroVadDataset(config, mode='train')
+    train_loader = torch.utils.data.DataLoader(train_dataset,
+                                               batch_size=config.batch_size,
+                                               collate_fn=SileroVadPadder,
+                                               num_workers=config.num_workers)
+
+    val_dataset = SileroVadDataset(config, mode='val')
+    val_loader = torch.utils.data.DataLoader(val_dataset,
+                                             batch_size=config.batch_size,
+                                             collate_fn=SileroVadPadder,
+                                             num_workers=config.num_workers)
+
+    if config.jit_model_path:
+        print(f'Loading model from the local folder: {config.jit_model_path}')
+        model = init_jit_model(config.jit_model_path, device=config.device)
+    else:
+        if config.use_torchhub:
+            print('Loading model using torch.hub')
+            model, _ = torch.hub.load(repo_or_dir='snakers4/silero-vad',
+                                      model='silero_vad',
+                                      onnx=False,
+                                      force_reload=True)
+        else:
+            print('Loading model using silero-vad library')
+            from silero_vad import load_silero_vad
+            model = load_silero_vad(onnx=False)
+
+    print('Model loaded')
+    model.to(config.device)
+    decoder = VADDecoderRNNJIT().to(config.device)
+    decoder.load_state_dict(model._model_8k.decoder.state_dict() if config.tune_8k else model._model.decoder.state_dict())
+    decoder.train()
+    params = decoder.parameters()
+    optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, params),
+                                 lr=config.learning_rate)
+    criterion = nn.BCELoss(reduction='none')
+
+    best_val_roc = 0
+    for i in range(config.num_epochs):
+        print(f'Starting epoch {i + 1}')
+        train_loss = train(config, train_loader, model, decoder, criterion, optimizer, config.device)
+        val_loss, val_roc = validate(config, val_loader, model, decoder, criterion, config.device)
+        print(f'Metrics after epoch {i + 1}:\n'
+              f'\tTrain loss: {round(train_loss, 3)}\n',
+              f'\tValidation loss: {round(val_loss, 3)}\n'
+              f'\tValidation ROC-AUC: {round(val_roc, 3)}')
+
+        if val_roc > best_val_roc:
+            print('New best ROC-AUC, saving model')
+            best_val_roc = val_roc
+            if config.tune_8k:
+                model._model_8k.decoder.load_state_dict(decoder.state_dict())
+            else:
+                model._model.decoder.load_state_dict(decoder.state_dict())
+            torch.jit.save(model, config.model_save_path)
+    print('Done')

tuning/utils.py

@@ -0,0 +1,357 @@
+from sklearn.metrics import roc_auc_score, accuracy_score
+from torch.utils.data import Dataset
+import torch.nn as nn
+from tqdm import tqdm
+import pandas as pd
+import numpy as np
+import torchaudio
+import warnings
+import random
+import torch
+import gc
+warnings.filterwarnings('ignore')
+
+
+def read_audio(path: str,
+               sampling_rate: int = 16000,
+               normalize=False):
+
+    wav, sr = torchaudio.load(path)
+
+    if wav.size(0) > 1:
+        wav = wav.mean(dim=0, keepdim=True)
+
+    if sampling_rate:
+        if sr != sampling_rate:
+            transform = torchaudio.transforms.Resample(orig_freq=sr,
+                                                       new_freq=sampling_rate)
+            wav = transform(wav)
+            sr = sampling_rate
+
+    if normalize and wav.abs().max() != 0:
+        wav = wav / wav.abs().max()
+
+    return wav.squeeze(0)
+
+
+def build_audiomentations_augs(p):
+    from audiomentations import SomeOf, AirAbsorption, BandPassFilter, BandStopFilter, ClippingDistortion, HighPassFilter, HighShelfFilter, \
+                                LowPassFilter, LowShelfFilter, Mp3Compression, PeakingFilter, PitchShift, RoomSimulator, SevenBandParametricEQ, \
+                                Aliasing, AddGaussianNoise
+    transforms = [Aliasing(p=1),
+                  AddGaussianNoise(p=1),
+                  AirAbsorption(p=1),
+                  BandPassFilter(p=1),
+                  BandStopFilter(p=1),
+                  ClippingDistortion(p=1),
+                  HighPassFilter(p=1),
+                  HighShelfFilter(p=1),
+                  LowPassFilter(p=1),
+                  LowShelfFilter(p=1),
+                  Mp3Compression(p=1),
+                  PeakingFilter(p=1),
+                  PitchShift(p=1),
+                  RoomSimulator(p=1, leave_length_unchanged=True),
+                  SevenBandParametricEQ(p=1)]
+    tr = SomeOf((1, 3), transforms=transforms, p=p)
+    return tr
+
+
+class SileroVadDataset(Dataset):
+    def __init__(self,
+                 config,
+                 mode='train'):
+
+        self.num_samples = 512  # constant, do not change
+        self.sr = 16000  # constant, do not change
+
+        self.resample_to_8k = config.tune_8k
+        self.noise_loss = config.noise_loss
+        self.max_train_length_sec = config.max_train_length_sec
+        self.max_train_length_samples = config.max_train_length_sec * self.sr
+
+        assert self.max_train_length_samples % self.num_samples == 0
+        assert mode in ['train', 'val']
+
+        dataset_path = config.train_dataset_path if mode == 'train' else config.val_dataset_path
+        self.dataframe = pd.read_feather(dataset_path).reset_index(drop=True)
+        self.index_dict = self.dataframe.to_dict('index')
+        self.mode = mode
+        print(f'DATASET SIZE : {len(self.dataframe)}')
+
+        if mode == 'train':
+            self.augs = build_audiomentations_augs(p=config.aug_prob)
+        else:
+            self.augs = None
+
+    def __getitem__(self, idx):
+        idx = None if self.mode == 'train' else idx
+        wav, gt, mask = self.load_speech_sample(idx)
+
+        if self.mode == 'train':
+            wav = self.add_augs(wav)
+            if len(wav) > self.max_train_length_samples:
+                wav = wav[:self.max_train_length_samples]
+                gt = gt[:int(self.max_train_length_samples / self.num_samples)]
+                mask = mask[:int(self.max_train_length_samples / self.num_samples)]
+
+        wav = torch.FloatTensor(wav)
+        if self.resample_to_8k:
+            transform = torchaudio.transforms.Resample(orig_freq=self.sr,
+                                                       new_freq=8000)
+            wav = transform(wav)
+        return wav, torch.FloatTensor(gt), torch.from_numpy(mask)
+
+    def __len__(self):
+        return len(self.index_dict)
+
+    def load_speech_sample(self, idx=None):
+        if idx is None:
+            idx = random.randint(0, len(self.index_dict) - 1)
+        wav = read_audio(self.index_dict[idx]['audio_path'], self.sr).numpy()
+
+        if len(wav) % self.num_samples != 0:
+            pad_num = self.num_samples - (len(wav) % (self.num_samples))
+            wav = np.pad(wav, (0, pad_num), 'constant', constant_values=0)
+
+        gt, mask = self.get_ground_truth_annotated(self.index_dict[idx]['speech_ts'], len(wav))
+
+        assert len(gt) == len(wav) / self.num_samples
+
+        mask[gt == 0]
+
+        return wav, gt, mask
+
+    def get_ground_truth_annotated(self, annotation, audio_length_samples):
+        gt = np.zeros(audio_length_samples)
+
+        for i in annotation:
+            gt[int(i['start'] * self.sr): int(i['end'] * self.sr)] = 1
+
+        squeezed_predicts = np.average(gt.reshape(-1, self.num_samples), axis=1)
+        squeezed_predicts = (squeezed_predicts > 0.5).astype(int)
+        mask = np.ones(len(squeezed_predicts))
+        mask[squeezed_predicts == 0] = self.noise_loss
+        return squeezed_predicts, mask
+
+    def add_augs(self, wav):
+        while True:
+            try:
+                wav_aug = self.augs(wav, self.sr)
+                if np.isnan(wav_aug.max()) or np.isnan(wav_aug.min()):
+                    return wav
+                return wav_aug
+            except Exception as e:
+                continue
+
+
+def SileroVadPadder(batch):
+    wavs = [batch[i][0] for i in range(len(batch))]
+    labels = [batch[i][1] for i in range(len(batch))]
+    masks = [batch[i][2] for i in range(len(batch))]
+
+    wavs = torch.nn.utils.rnn.pad_sequence(
+        wavs, batch_first=True, padding_value=0)
+
+    labels = torch.nn.utils.rnn.pad_sequence(
+        labels, batch_first=True, padding_value=0)
+
+    masks = torch.nn.utils.rnn.pad_sequence(
+        masks, batch_first=True, padding_value=0)
+
+    return wavs, labels, masks
+
+
+class VADDecoderRNNJIT(nn.Module):
+
+    def __init__(self):
+        super(VADDecoderRNNJIT, self).__init__()
+
+        self.rnn = nn.LSTMCell(128, 128)
+        self.decoder = nn.Sequential(nn.Dropout(0.1),
+                                     nn.ReLU(),
+                                     nn.Conv1d(128, 1, kernel_size=1),
+                                     nn.Sigmoid())
+
+    def forward(self, x, state=torch.zeros(0)):
+        x = x.squeeze(-1)
+        if len(state):
+            h, c = self.rnn(x, (state[0], state[1]))
+        else:
+            h, c = self.rnn(x)
+
+        x = h.unsqueeze(-1).float()
+        state = torch.stack([h, c])
+        x = self.decoder(x)
+        return x, state
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+def train(config,
+          loader,
+          jit_model,
+          decoder,
+          criterion,
+          optimizer,
+          device):
+
+    losses = AverageMeter()
+    decoder.train()
+
+    context_size = 32 if config.tune_8k else 64
+    num_samples = 256 if config.tune_8k else 512
+    stft_layer = jit_model._model_8k.stft if config.tune_8k else jit_model._model.stft
+    encoder_layer = jit_model._model_8k.encoder if config.tune_8k else jit_model._model.encoder
+
+    with torch.enable_grad():
+        for _, (x, targets, masks) in tqdm(enumerate(loader), total=len(loader)):
+            targets = targets.to(device)
+            x = x.to(device)
+            masks = masks.to(device)
+            x = torch.nn.functional.pad(x, (context_size, 0))
+
+            outs = []
+            state = torch.zeros(0)
+            for i in range(context_size, x.shape[1], num_samples):
+                input_ = x[:, i-context_size:i+num_samples]
+                out = stft_layer(input_)
+                out = encoder_layer(out)
+                out, state = decoder(out, state)
+                outs.append(out)
+            stacked = torch.cat(outs, dim=2).squeeze(1)
+
+            loss = criterion(stacked, targets)
+            loss = (loss * masks).mean()
+            loss.backward()
+            optimizer.step()
+            losses.update(loss.item(), masks.numel())
+
+    torch.cuda.empty_cache()
+    gc.collect()
+
+    return losses.avg
+
+
+def validate(config,
+             loader,
+             jit_model,
+             decoder,
+             criterion,
+             device):
+
+    losses = AverageMeter()
+    decoder.eval()
+
+    predicts = []
+    gts = []
+
+    context_size = 32 if config.tune_8k else 64
+    num_samples = 256 if config.tune_8k else 512
+    stft_layer = jit_model._model_8k.stft if config.tune_8k else jit_model._model.stft
+    encoder_layer = jit_model._model_8k.encoder if config.tune_8k else jit_model._model.encoder
+
+    with torch.no_grad():
+        for _, (x, targets, masks) in tqdm(enumerate(loader), total=len(loader)):
+            targets = targets.to(device)
+            x = x.to(device)
+            masks = masks.to(device)
+            x = torch.nn.functional.pad(x, (context_size, 0))
+
+            outs = []
+            state = torch.zeros(0)
+            for i in range(context_size, x.shape[1], num_samples):
+                input_ = x[:, i-context_size:i+num_samples]
+                out = stft_layer(input_)
+                out = encoder_layer(out)
+                out, state = decoder(out, state)
+                outs.append(out)
+            stacked = torch.cat(outs, dim=2).squeeze(1)
+
+            predicts.extend(stacked[masks != 0].tolist())
+            gts.extend(targets[masks != 0].tolist())
+
+            loss = criterion(stacked, targets)
+            loss = (loss * masks).mean()
+            losses.update(loss.item(), masks.numel())
+    score = roc_auc_score(gts, predicts)
+
+    torch.cuda.empty_cache()
+    gc.collect()
+
+    return losses.avg, round(score, 3)
+
+
+def init_jit_model(model_path: str,
+                   device=torch.device('cpu')):
+    torch.set_grad_enabled(False)
+    model = torch.jit.load(model_path, map_location=device)
+    model.eval()
+    return model
+
+
+def predict(model, loader, device, sr):
+    with torch.no_grad():
+        all_predicts = []
+        all_gts = []
+        for _, (x, targets, masks) in tqdm(enumerate(loader), total=len(loader)):
+            x = x.to(device)
+            out = model.audio_forward(x, sr=sr)
+
+            for i, out_chunk in enumerate(out):
+                predict = out_chunk[masks[i] != 0].cpu().tolist()
+                gt = targets[i, masks[i] != 0].cpu().tolist()
+
+                all_predicts.append(predict)
+                all_gts.append(gt)
+    return all_predicts, all_gts
+
+
+def calculate_best_thresholds(all_predicts, all_gts):
+    best_acc = 0
+    for ths_enter in tqdm(np.linspace(0, 1, 20)):
+        for ths_exit in np.linspace(0, 1, 20):
+            if ths_exit >= ths_enter:
+                continue
+
+            accs = []
+            for j, predict in enumerate(all_predicts):
+                predict_bool = []
+                is_speech = False
+                for i in predict:
+                    if i >= ths_enter:
+                        is_speech = True
+                        predict_bool.append(1)
+                    elif i <= ths_exit:
+                        is_speech = False
+                        predict_bool.append(0)
+                    else:
+                        val = 1 if is_speech else 0
+                        predict_bool.append(val)
+
+                score = round(accuracy_score(all_gts[j], predict_bool), 4)
+                accs.append(score)
+
+            mean_acc = round(np.mean(accs), 3)
+            if mean_acc > best_acc:
+                best_acc = mean_acc
+                best_ths_enter = round(ths_enter, 2)
+                best_ths_exit = round(ths_exit, 2)
+    return best_ths_enter, best_ths_exit, best_acc