Распознавание речи с Wav2Letter

В этом туториале мы рассмотрим, как с помощью библиотеки PuzzleLib можно построить систему распознавания речи. Мы обучим на открытом наборе данных LibriSpeech уже собранную на PuzzleLib нейронную сеть Wav2Letter.

Подготовка данных

Для того, чтобы запустить обучение модели нужно подготовить специальные csv файлы (манифесты) для обучающих, валидационных и тестовых данных. Строчка манифест-файла содержит путь к аудиофайлу и расшифровку произносимой в аудиофайле речи.

Скрипт PrepareLibrispeech.py позволяет скачать набор данных LibriSpeech, сконвертировать аудиофайлы из flac в wav и создать 6 манифест-файлов: по два (чистые и более шумные данные) для обучения, валидации и тестирования. Строчки манифест-файлов упорядочены по длинам соотвествующих аудиозаписей. Это способствует более эффективному обучению модели. Аргументами скрипта являются полный путь для сохранения данных dataDir и путь для сохранения манифест-файлов manifestDir.

Извлечение признаков

Существуют разные способы извлечения признаков из аудиоданных, мы используем нормализованные спектрограммы, так как они позволяют достаточно полно описать данные. Гиперпараметры, а также словарь токенов задаются в классе LibriConfig файла Config.py. Словарь токенов labels состоит из латинского алфавита, апострофа, пробела и символа пустышки, который соответсвует пропуску буквы.

class LibriConfig:
    sampleRate = 16000
    windowSize = 0.02
    windowStride = 0.01
    window = 'hamming'
    labels = "_'abcdefghijklmnopqrstuvwxyz "

Применяемые к аудио преобразования реализованы в функции preprocess скрипта Data/dataLoader.py. Аудио считывается с определенной частотой дискретизации sampleRate, затем к сигналу применяется оконное преобразование Фурье с заданным размером окна windowSize, шагом windowStride и типом оконной функции window. Из полученной комплекснозначной спектрограммы извлекается ее вещественная часть. Описанные методы реализованы в библиотеке librosa.

def preprocess(audioPath, sampleRate, windowSize, windowStride, window):
    y = loadAudio(audioPath)

    nFft = int(sampleRate * windowSize)
    winLength = nFft
    hopLength = int(sampleRate * windowStride)

    D = librosa.stft(y, n_fft=nFft, hop_length=hopLength, win_length=winLength, window=window)

    spect, phase = librosa.magphase(D)
И наконец к спектрограмме применяется per-channel energy normalization, реализованная в функции pcen. Это популярная в распознавании речи технология нормализации аудио, позволяющая ослабить фоновой шум и подчеркнуть звук переднего плана.
    pcenResult = pcen(E=spect, sr=sampleRate, hopLength=hopLength)
    return pcenResult
Класс SpectrogramDataset инициализируется манифест-файлом и конфигом. Он описывает датасет аудиоданных с расшифровками, получая пути к аудиофайлам и траскрипции из манифест-файла, а также создает словарь токенов labelsMap и задает гиперпараметры обработки аудиоданных.
class SpectrogramDataset(object):
    def __init__(self, manifestFilePath, config):
        self.manifestFilePath = manifestFilePath
        with open(self.manifestFilePath) as f:
            ids = f.readlines()
        self.ids = [x.strip().split(',') for x in ids]
        self.size = len(ids)
        self.labelsMap = dict([(config.labels[i], i) for i in range(len(config.labels))])
        self.sampleRate = config.sampleRate
        self.windowSize = config.windowSize
        self.windowStride = config.windowStride
        self.window = config.window
Метод __getitem__ возвращает вычисленную с помощью функции preprocess спектрограмму, переведенную в последовательность индексов токенов расшифровку, путь к аудио файлу и саму текстовую расшифровку.
    def __getitem__(self, index):
        sample = self.ids[index]
        audioPath, transcriptLoaded = sample[0], sample[1]
        spect = preprocess(audioPath, self.sampleRate, self.windowSize, self.windowStride, self.window)

        transcript = list(filter(None, [self.labelsMap.get(x) for x in list(transcriptLoaded)]))
        return spect, transcript, audioPath, transcriptLoaded
Для агрегации данных в батчи реализован класс BucketingSampler. Он при необходимости перемешивает индексы данных, а так же группирует их в бины, размер бина соответствует размеру батча.
class BucketingSampler(object):
    def __init__(self, dataSource, batchSize=1, shuffle=False):    
        self.dataSource = dataSource
        self.batchSize = batchSize
        self.ids = list(range(0, len(dataSource)))
        self.shuffle = shuffle
        self.reset()

    ...

    def getBins(self):
        if self.shuffle:
            np.random.shuffle(self.ids)
        self.bins = [self.ids[i:i + self.batchSize] for i in range(0, len(self.ids), self.batchSize)]


    def reset(self):
        self.getBins()
        self.batchId = 0
Класс DataLoader, используя эти группы индексов, позволяет итеративно генерировать батчи.
class DataLoader(object):        
    def __init__(self, dataset, batchSampler):
        self.dataset = dataset
        self.batchSampler = batchSampler
        self.sampleIter = iter(self.batchSampler)


    def __next__(self):
        try:
            indices = next(self.sampleIter)
            indices = [i for i in indices][0]
            batch = getBatch([self.dataset[i] for i in indices])        
            return batch
        except:
            raise StopIteration()
Для формирования батчей используется функция getBatch. Извлеченные из аудио сигналов признаки приводятся к одной максимальной длине с помощью заворачивающего паддинга и агрегируются в трехмерный тензор inputs. Поэтому мы упорядочили аудио в манифест-файлах по длительности, так как эффективно агрегировать в один батч аудио сигналы по длине. Относительные длины сигналов записываются в inputPercentages.
        seqLength = tensor.shape[1]
        tensorNew = np.pad(tensor, ((0, 0), (0, abs(seqLength-maxSeqlength))), 'wrap')
        inputs[x] = tensorNew
        inputPercentages[x] = seqLength / float(maxSeqlength)
Последовательности индексов токенов расшифровок комбинируются в одну targets, а длины исходных последовательностей записываются в targetSizes. Указанные переменные вместе с путем к аудиофайлу и текстом расшифровки формируют батч.
        targetSizes[x] = len(target)
        targets.extend(target)
        inputFilePathAndTranscription.append([tensorPath, orignalTranscription])

    targets = np.array(targets)
    return inputs, inputPercentages, targets, targetSizes, inputFilePathAndTranscription

Система распознавания речи

Система распознавания речи состоит из двух модулей: акустическая модель и декодинг. Акустическая модель устанавливает зависимость между аудио сигналом и распределением вероятностей токенов на нем. В нашем случае это сверточная нейронная сеть, принимающая на вход последовательность признаков, извлеченных из аудио сигнала, и выдающая последовательность векторов вероятностей токенов на нем. То есть модель предсказывает вероятности токенов на каждом фиксированном временном окне. Вот, что получится, если взять из выхода акустической модели токены с максимальной вероятностью: ______________п_р__и_в__е__тт__ _м___и__ррр__. Для того, чтобы преобразовать последовательность векторов вероятностей токенов в текст привет мир выполняется декодинг.

Обучение акустической модели и валидация на соответсвующих манифест-файлах реализованы в скрипте Train.py. Аргументами скрипта являются:

  • пути к обучающему и валидационным манифест-файлам, если передано None, то обучение или валидация будут пропущены
    parser.add_argument('--trainManifest', metavar='DIR',
                        help='path to train manifest csv', default='Data/train-clean.csv')
    parser.add_argument('--valManifests', metavar='DIR',
                        help='path to validation manifest csv', default=['Data/dev-clean.csv'])
    
  • размер батча, количество эпох обучения, learning rate
    parser.add_argument('--batchSize', default=8, type=int, help='Batch size for training')
    parser.add_argument('--epochs', default=100, type=int, help='Number of training epochs')
    parser.add_argument('--lr', '--learningRate', default=1e-5, type=float, help='initial learning rate')
    
  • аргументы отвечающие за то, как часто сохранять веса акустической модели, префикс названий файлов весов, путь к директории, куда сохранять веса
    parser.add_argument('--checkpointPerBatch', default=3000, type=int, help='Save checkpoint per batch. 0 means never save')
    parser.add_argument('--checkpointName', default='w2l', type=str, help='Name of checkpoints')
    parser.add_argument('--saveFolder', default='Checkpoints/', help='Location to save epoch models')
    
  • путь к весам акустической модели, которыми ее инициализировать
    parser.add_argument('--continueFrom', default=None, help='Continue from checkpoint model')
    
    С помощью описанных ранее классов работы с данными и функции getDataLoader инициализируется загрузчики обучающих и валидационных данных: trainLoader и valLoaders.
    from Data.dataLoader import DataLoader, SpectrogramDataset, BucketingSampler
    ...
    def getDataLoader(manifestFilePath, config, batchSize):
        dataset = SpectrogramDataset(manifestFilePath=manifestFilePath, config=config)
        sampler = BucketingSampler(dataset, batchSize=batchSize)
        return DataLoader(dataset, batchSampler=sampler)
    

Обучение акустической модели

Используемая нами архитектура нейронной сети акустической модели вдохновлена технологией Wav2Letter. Преимущество такой акустической модели в том, что она полностью состоит из сверточных слоев, а значит вычисляется эффективно. Сеть Wav2Letter уже собрана на PuzzleLib, поэтому ее просто нужно импортировать из библиотеки.

from PuzzleLib.Models.Nets.WaveToLetter import loadW2L
Размерность входа нейронной сети определяется размерностью извлеченных признаков, а именно размером окна, использовавшегося при вычислении оконного преобразование Фурье для аудиосигнала. Размерность выходных данных соотвествуется количеству токенов. Если подан аргумент continueFrom, то модель инициализируется сохраннеными весами.
    nfft = int(config.sampleRate * config.windowSize)
    w2l = loadW2L(modelpath=args.continueFrom, inmaps=(1 + nfft // 2), nlabels=len(config.labels))
Большинство наборов данных для распознавания речи содержат только аудиофайлы и их расшифровки, но не содержат информацию о том, каким временным сегментам соотвествуют какие токены расшифровки. К тому же несколько последовательностей токенов могут соответсвовать одной и той же расшифровке. Например, последовательности токенов для "привет мир":
______________п_р__и_в__е__тт__ _м___и__ррр__
______________п_р____иии_в__еее__т__ ____м___и__р__
___________ппп__р__ии____ввв__е__тт__ _ммм_и_р__
Поэтому для обучения нейросети используется технология Connectionist Temporal Classification. CTC максимизирует вероятности всех возможных последовательностей токенов, которые при данной длине последовательности соответсвуют данной расшифровке. CTC критерий так же реализован в PuzzleLib. Для его инициализации лишь нужно подать индекс токена-пустышки.
from PuzzleLib.Cost.CTC import CTC

...

    blankIndex = [i for i in range(len(config.labels)) if config.labels[i] == '_'][0]
    ctc = CTC(blankIndex)
В качестве метода оптимизации используется Adam.
from PuzzleLib.Optimizers.Adam import Adam

...

    adam = Adam(alpha=args.lr)
    adam.setupOn(w2l, useGlobalState=True)
Для обучения модели вызывается функция train. Модель переводится в режим обучения, затем загрузжчик данных итеративно генерирует батчи. Батч переносится на видеокарту и к нему применяется модель.
def train(model, ctc, optimizer, loader, checkpointPerBatch, saveFolder, saveName):
    model.reset()
    model.trainMode()
    loader.reset()

    for i, (data) in enumerate(loader):
        inputs, inputPercentages, targets, targetSizes, _ = data
        gpuInput = gpuarray.to_gpu(inputs.astype(np.float32))
        out = model(gpuInput)
Длины предсказанных последовательностей outlen оценениваются через относительные длины входных признаков.
        outlen = gpuarray.to_gpu((out.shape[0] * inputPercentages).astype(np.int32))
По предсказаниям модели out, длинам предсказаний outlen, расшифровкам targets и их длинам targetSizes вычисляется функция потерь CTC и градиент.
        error, grad = ctc([out, outlen], [targets, targetSizes])
Параметры модели обновляются с помощью метода обратного распространения ошибки.

        optimizer.zeroGradParams()
        model.backward(grad, updGrad=False)
        optimizer.update()

Декодинг и валидация

Выходом акустической модели является тензор размерности (N_{batch}, T, N_{tokens}). Стоит задача получить из последовательности векторов вероятностей токенов текст. Мы рассмотрим самый простой подход - использование Greedy Decoder. Он реализован в файле Decoder.py. Для его инициализации нужно подать последовательность токенов и индекс символа пустышки.

from Decoder import GreedyDecoder

...

    decoder = GreedyDecoder(config.labels, blankIndex) 
Построение текста предсказания реализовано в методе decode. На каждом временном шаге декодер использует токен с максимальной вероятностью. Например, получаем последовательность индексов токенов, соответсвующую последовательности токенов ___________ппп__р__ии____ввв__е__тт__ _ммм_и_р__.
    def decode(self, probs, sizes=None):
        npMaxProbs = np.argmax(probs, 2)
Данная последовательность индексов токенов обрабатывается с помощью метода processString. Итеративно проходясь по последовательности, новый символ char добавляется к строке предсказания string, только если он не является символом пустышкой и не совпадает с предыдущим символом в последовательности. Таким образом, мы получаем текст привет мир
    def processString(self, sequence, size):
        string = ''
        for i in range(size):
            char = self.intToChar[sequence[i].item()]
            if char != self.intToChar[self.blankIndex]:
                if i != 0 and char == self.intToChar[sequence[i - 1].item()]:
                    pass
                elif char == self.labels[self.spaceIndex]:
                    string += ' '
                else:
                    string = string + char
        return string  
Чтобы сравнить предсказанный текст с исходной расшифровкой используются метрики Word Error Rate и Character Error Rate. Обе метрики основаны на вычислении расстояния Левенштейна. Эта метрика измеряет разность между двумя последовательностями. Она определяется как минимальное количество одноэлементных операций (а именно вставки, удаления, замены), необходимых для превращения одной последовательности символов в другую. Метрика реализована в библиотеке Levenshtein. Для подсчета WER между двумя текстами s1 и s2 реализована функция calcWer. Тексты представляются как последовательности слов. Затем вычисляется расстояние Левенштейна между двумя последовательностями и нормируется длиной второй последовательности.
def calcWer(s1, s2):
    b = set(s1.split() + s2.split())
    word2char = dict(zip(b, range(len(b))))
    w1 = [chr(word2char[w]) for w in s1.split()]
    w2 = [chr(word2char[w]) for w in s2.split()]
    return Levenshtein.distance(''.join(w1), ''.join(w2)) / len(''.join(w2))
Для подсчета CER реализована функция calcCer. Тексты представляются как последовательности символов. Затем между ними вычисляется расстояние Левенштейна и нормируется длиной второй последовательности.
def calcCer(s1, s2):
    s1, s2, = s1.replace(' ', ''), s2.replace(' ', '')
    return Levenshtein.distance(s1, s2) / len(s2)
Валидация реализована в функции validate. Модель переводится в режим применения. Метрики качества totalCer и totalWer инициализируются нулями.
def validate(model, loader, decoder, logPath):
    loader.reset()
    model.evalMode()
    totalCer, totalWer = 0, 0
Как и при обучении к итеративно сгенерированным батчам применяется модель, оцениваются длины предсказанных последовательностей.
    for i, (data) in enumerate(loader):
        inputs, inputPercentages, targets, targetSizes, inputFile = data
        gpuInput = gpuarray.to_gpu(inputs.astype(np.float32))
        out = model(gpuInput)
        outlen = (out.shape[0] * inputPercentages).astype(np.int32)
Для получения текстовых предсказаний используется метод decode.
        decodedOutput = decoder.decode(np.moveaxis(out.get(), 0, 1), outlen)
Для каждого предсказания из батча вычисляется wer и cer. Вычисленные на предсказании метрики добавляются к общим totalCer и totalWer.
        wer, cer = 0, 0
        for x in range(len(decodedOutput)):            
            transcript, reference = decodedOutput[x], inputFile[x][1]
            print ('transcript: {}\nreference: {}\nfilepath: {}'.format(transcript, reference, inputFile[x][0]))
            try:
                wer += calcWer(transcript, reference)
                cer += calcCer(transcript, reference)
            except Exception as e:
                print ('encountered exception {}'.format(e))
        totalCer += cer
        totalWer += wer
После завершения цикла по батчам вычисляются общие wer и cer, нормированные на длину валидационного набора данных.
    wer = totalWer / len(loader.dataset) * 100
    cer = totalCer / len(loader.dataset) * 100

Результаты

После 42 эпох обучения на чистых обучающих данных удалось достигнуть следующего качества на чистых валидационных и тестовых данных LibriSpeech:

Dataset WER CER
dev-clean 17.526 6.611
test-clean 16.495 6.131