NVIDIAのNeMoを使って自動音声認識モデルの学習、最適化、解析、可視化、デプロイを行う

# Download the dataset
print("******")
if not os.path.exists(data_dir + '/an4_sphere.tar.gz'):
    an4_url = 'http://www.speech.cs.cmu.edu/databases/an4/an4_sphere.tar.gz'
    an4_path = wget.download(an4_url, data_dir)
    print(f"Dataset downloaded at: {an4_path}")
else:
    print(f"Tarfile already exists at {data_dir + '/an4_sphere.tar.gz'}")
    an4_path = data_dir + '/an4_sphere.tar.gz'

# convert .sph to .wav
if not os.path.exists(data_dir + '/an4/'):
    # Untar and convert .sph to .wav (using sox)
    tar = tarfile.open(an4_path)
    tar.extractall(path=data_dir)
    sph_list = glob.glob(data_dir + '/an4/**/*.sph', recursive=True)
    for sph_path in tqdm(sph_list):
        wav_path = sph_path[:-4] + '.wav'
        cmd = ["sox", sph_path, wav_path]
        subprocess.run(cmd)


import librosa
import IPython.display as ipd

# Load and listen to the audio file
example_file = data_dir + '/an4/wav/an4_clstk/mgah/cen2-mgah-b.wav'
audio, sample_rate = librosa.load(example_file)

ipd.Audio(example_file, rate=sample_rate)
_ = librosa.display.waveplot(audio)

# Plot the mel spectrogram of our sample
mel_spec = librosa.feature.melspectrogram(audio, sr=sample_rate)
mel_spec_db = librosa.power_to_db(mel_spec, ref=np.max)
librosa.display.specshow(
    mel_spec_db, x_axis='time', y_axis='mel')
plt.colorbar()
plt.title('Mel Spectrogram');

# NeMo's "core" package
import nemo
# NeMo's ASR collection - this collections contains complete ASR models and
# building blocks (modules) for ASR
import nemo.collections.asr as nemo_asr
# This line will download pre-trained QuartzNet15x5 model from NVIDIA's NGC cloud and instantiate it for you
quartznet = nemo_asr.models.EncDecCTCModel.from_pretrained(model_name="QuartzNet15x5Base-En")#

files = ['./an4/wav/an4_clstk/mgah/cen2-mgah-b.wav']
for fname, transcription in zip(files, quartznet.transcribe(paths2audio_files=files)):
  print(f"Audio in {fname} was recognized as: {transcription}")

{"audio_filepath": "path/to/audio.wav", "duration": 3.45, "text": "this is a nemo tutorial"}

...
<s> P I T T S B U R G H </s> (cen5-fash-b)
<s> TWO SIX EIGHT FOUR FOUR ONE EIGHT </s> (cen7-fash-b)
...

# Function to build a manifest
def build_manifest(transcripts_path, manifest_path, wav_path):
    audio_paths=[]; durations=[]; texts=[]
    with open(transcripts_path, 'r') as fin:
        with open(manifest_path, 'w') as fout:
            for line in fin:
                # Lines look like this:
                # <s> transcript </s> (fileID)
                transcript = line[: line.find('(')-1].lower()
                transcript = transcript.replace('<s>', '').replace('</s>', '')
                transcript = transcript.strip()

                file_id = line[line.find('(')+1 : -2]  # e.g. "cen4-fash-b"
                audio_path = os.path.join(
                    data_dir, wav_path,
                    file_id[file_id.find('-')+1 : file_id.rfind('-')],
                    file_id + '.wav')

                duration = librosa.core.get_duration(filename=audio_path)

                audio_paths.append(audio_path)
                durations.append(duration)
                texts.append(transcript)
                # Write the metadata to the manifest
                metadata = {
                    "audio_filepath": audio_path,
                    "duration": duration,
                    "text": transcript
                }
                json.dump(metadata, fout)
                fout.write('\n') 
    return audio_paths, durations, texts


# Building Manifests
train_transcripts = data_dir + '/an4/etc/an4_train.transcription'
train_manifest = data_dir + '/an4/train_manifest.json'
train_audio_paths, train_durations, train_texts = build_manifest(train_transcripts, train_manifest, 'an4/wav/an4_clstk')

test_transcripts = data_dir + '/an4/etc/an4_test.transcription'
test_manifest = data_dir + '/an4/test_manifest.json'
test_audio_paths, test_durations, test_texts = build_manifest(test_transcripts, test_manifest, 'an4/wav/an4test_clstk')

{"audio_filepath": "./an4/wav/an4_clstk/fash/an251-fash-b.wav", "duration": 1.0, "text": "yes"}
{"audio_filepath": "./an4/wav/an4_clstk/fash/an253-fash-b.wav", "duration": 0.7, "text": "go"}
{"audio_filepath": "./an4/wav/an4_clstk/fash/an254-fash-b.wav", "duration": 0.9, "text": "yes"}
{"audio_filepath": "./an4/wav/an4_clstk/fash/an255-fash-b.wav", "duration": 2.6, "text": "u m n y h six"}
{"audio_filepath": "./an4/wav/an4_clstk/fash/cen1-fash-b.wav", "duration": 3.5, "text": "h i n i c h"}

def save_spectogram_as_img(audio_path, datadir, plt_type='spec'):
    filename = os.path.basename(audio_path)
    out_path = os.path.join(datadir, filename.replace('.wav', '.png'))
    audio, sample_rate = librosa.load(audio_path)
    if plt_type=='spec':
        spec = np.abs(librosa.stft(audio))
        spec_db = librosa.amplitude_to_db(spec, ref=np.max)
    else:
        mel_spec = librosa.feature.melspectrogram(audio, sr=sample_rate)
        mel_spec_db = librosa.power_to_db(mel_spec, ref=np.max)
    
    fig = plt.Figure()
    ax = fig.add_subplot()
    ax.set_axis_off()
    
    librosa.display.specshow(
        spec_db if plt_type=='spec' else mel_spec_db, 
        y_axis='log' if plt_type=='spec' else 'mel', 
        x_axis='time', ax=ax)

    fig.savefig(out_path)

# convert audio file to spectogram and mel spectogram images
if not os.path.exists('./an4/melspectogram_images/'):
    for path in tqdm(train_audio_paths): 
        save_mel_spectogram_as_img(path, datadir='./an4/images/')
        save_spectogram_as_img(path, datadir='./an4/melspectogram_images/', plt_type='mel')

# log filename, playable audio, duration of audio, transcript, spectogram and mel spectogram to W&B for ease of reference
if LOG_WANDB:
    # create W&B Table
    wandb.init(project="ASR")

    audio_table = wandb.Table(
        columns=['Filename', 'Audio File', 'Duration', 'Transcript', 'Spectogram', 'Mel-Spectogram'])

    for path, duration, text in zip(train_audio_paths, train_durations, train_texts):
        filename = os.path.basename(path)
        img_fn   = filename.replace('.wav', '.png')
        spec_pth = os.path.join('./an4/images', img_fn)
        melspec_pth = os.path.join('./an4/melspectogram_images', img_fn)
        audio_table.add_data(
            filename, wandb.Audio(path), duration, text, wandb.Image(spec_pth), wandb.Image(melspec_pth))

    wandb.log({"Train Data": audio_table})
    wandb.finish();

- filters: 128
  repeat: 1
  kernel: [11]
  stride: [2]
  dilation: [1]
  dropout: 0.2
  residual: false
  separable: true
  se: true
  se_context_size: -1

import pytorch_lightning as pl
trainer = pl.Trainer(gpus=1, max_epochs=50)

import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger  

# initialize W&B logger and specify project name to store results to
wandb_logger = WandbLogger(project="ASR", log_model='all')

# set config params for W&B experiment
for k,v in params.items(): 
    wandb_logger.experiment.config[k]=v 

# initialize trainer with W&B logger    
trainer = pl.Trainer(gpus=1, max_epochs=10, logger=wandb_logger)

# Update train and test data path
params['model']['train_ds']['manifest_filepath'] = train_manifest
params['model']['validation_ds']['manifest_filepath'] = test_manifest

# initialize model
first_asr_model = nemo_asr.models.EncDecCTCModel(cfg=DictConfig(params['model']), trainer=trainer)

# Start training - this will automatically store results to Weights and Biases
trainer.fit(first_asr_model)
wandb.finish();

sweep_config = {
  "method": "random",   # Random search
  "metric": {           # We want to minimize `val_loss`
      "name": "val_loss",
      "goal": "minimize"
  },
  "parameters": {
        "lr": {
            # log uniform distribution between exp(min) and exp(max)
            "distribution": "log_uniform",
            "min": -9.21,   # exp(-9.21) = 1e-4
            "max": -4.61    # exp(-4.61) = 1e-2
        },
        "epoch": {
            "distribution": "int_uniform",
            "min": 3,
            "max": 10
        },
        "dropout": {
            "distribution": "uniform",
            "min": 0,  
            "max": 0.25     
        }
      
    }
}

    params['model']['optim']['lr'] = wandb.config.lr
    params['model']['encoder']['jasper'][-1]['dropout'] = wandb.config.dropout

def sweep_iteration():
    # load config
    config_path = './configs/config.yaml'
    yaml = YAML(typ='safe')
    with open(config_path) as f:
        params = yaml.load(f)    
        
    # set up W&B logger
    wandb.init()    # required to have access to `wandb.config`
    wandb_logger = WandbLogger(log_model='all')  # log final model

    # setup data
    params['model']['train_ds']['manifest_filepath'] = train_manifest
    params['model']['validation_ds']['manifest_filepath'] = test_manifest
    
    # setup sweep param
    params['model']['optim']['lr'] = wandb.config.lr
    params['model']['encoder']['jasper'][-1]['dropout'] = wandb.config.dropout
    
    trainer = pl.Trainer(gpus=1, max_epochs=wandb.config.epoch, logger=wandb_logger)
    
    # setup model - note how we refer to sweep parameters with wandb.config
    model = nemo_asr.models.EncDecCTCModel(cfg=DictConfig(params['model']), trainer=trainer)

    # train
    trainer.fit(model)

sweep_id = wandb.sweep(sweep_config, project="ASR")
wandb.agent(sweep_id, function=sweep_iteration)

artifact = run.use_artifact('user_name/project_name/new_artifact:v1', type='my_dataset') 
artifact_dir = artifact.download()

if LOG_WANDB:
    preds = quartznet.transcribe(paths2audio_files=test_audio_paths, batch_size=16)
    pred_table = wandb.Table(columns=['FileName', 'Audio', 'Prediction', 'Ground Truth'])

    for path, gt, pred in zip(test_audio_paths, test_texts, preds):
        pred_table.add_data(os.path.basename(path), wandb.Audio(path), pred, gt)

    run = wandb.init(project='ASR')
    wandb.log({'Prediction Table': pred_table})
    wandb.finish();

import copy
new_opt = copy.deepcopy(params['model']['optim'])
new_opt['lr'] = 0.001
first_asr_model.setup_optimization(optim_config=DictConfig(new_opt))

# Bigger batch-size = bigger throughput
params['model']['validation_ds']['batch_size'] = 16

# Setup the test data loader and make sure the model is on GPU
first_asr_model.setup_test_data(test_data_config=params['model']['validation_ds'])
first_asr_model.cuda()

# We will be computing Word Error Rate (WER) metric between our hypothesis and predictions.
# WER is computed as numerator/denominator.
# We'll gather all the test batches' numerators and denominators.
wer_nums = []
wer_denoms = []

# Loop over all test batches.
# Iterating over the model's `test_dataloader` will give us:
# (audio_signal, audio_signal_length, transcript_tokens, transcript_length)
# See the AudioToCharDataset for more details.
for test_batch in first_asr_model.test_dataloader():
        test_batch = [x.cuda() for x in test_batch]
        targets = test_batch[2]
        targets_lengths = test_batch[3]        
        log_probs, encoded_len, greedy_predictions = first_asr_model(
            input_signal=test_batch[0], input_signal_length=test_batch[1]
        )
        # Notice the model has a helper object to compute WER
        first_asr_model._wer.update(greedy_predictions, targets, targets_lengths)
        _, wer_num, wer_denom = first_asr_model._wer.compute()
        first_asr_model._wer.reset()
        wer_nums.append(wer_num.detach().cpu().numpy())
        wer_denoms.append(wer_denom.detach().cpu().numpy())

        # Release tensors from GPU memory
        del test_batch, log_probs, targets, targets_lengths, encoded_len, greedy_predictions

# We need to sum all numerators and denominators first. Then divide.
print(f"WER = {sum(wer_nums)/sum(wer_denoms)}")

# Check what kind of vocabulary/alphabet the model has right now
print(quartznet.decoder.vocabulary)

# Let's add "!" symbol there. Note that you can (and should!) change the vocabulary
# entirely when fine-tuning using a different language.
quartznet.change_vocabulary(
    new_vocabulary=[
        ' ', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n',
        'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', "'", "!"
    ]
)

>> [' ', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', "'"]
[NeMo I 2022-01-19 15:31:05 ctc_models:348] Changed decoder to output to [' ', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', "'", '!'] vocabulary.

# Use the smaller learning rate we set before
quartznet.setup_optimization(optim_config=DictConfig(new_opt))

# Point to the data we'll use for fine-tuning as the training set
quartznet.setup_training_data(train_data_config=params['model']['train_ds'])

# Point to the new validation data for fine-tuning
quartznet.setup_validation_data(val_data_config=params['model']['validation_ds'])

# And now we can create a PyTorch Lightning trainer and call `fit` again.
trainer = pl.Trainer(gpus=[1], max_epochs=2)
trainer.fit(quartznet)