# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
import numpy as np
import logging_util
import transducer_blocks
# set up logging
logger = logging_util.get_basic_logger(__name__)
class TranscriptionNetwork(transducer_blocks.Block):
""" Transcription Network (or Audio Encoding network) of the Transformer-Transducer model.
:param popart builder: popart builder object
:param int in_feats: input dimension
:param int subsampling_factor: sub-sampling factor for the initial convolutional subsampling layer
:param int enc_n_hid: encoder hidden dimension
:param int num_encoder_layers: the number of transformer layers for the transcription encoding network
:param int encoder_dim: dimension of the transformer layers
:param int num_attention_heads: number of attention heads
:param float enc_dropout: dropout rate for encoder net
:param int kernel_size: kernel size of the initial convolutional subsampling layer
"""
def __init__(self,
builder,
in_feats,
subsampling_factor,
num_encoder_layers,
encoder_dim,
num_attention_heads,
enc_dropout,
kernel_size=32,
min over modeler h max over critic f of E[((y-h(w))*f(x))**2]
"""
import argparse
import numpy as np
import time as time
import yaml
import tensorflow as tf
import tensorflow_probability as tfp
try:
from tensorflow.python import ipu
except:
pass
# set up logging
import logging_util
logger = logging_util.get_basic_logger("AdGMoM")
tfd = tfp.distributions
# function examples
def tau_fn_2dpoly(x):
return -1.5 * x + .9 * (x**2)
def tau_fn_3dpoly(x):
return -1.5 * x + .9 * (x**2) + x**3
def tau_fn_abs(x):
return np.abs(x)
import time
import os
import sys
import json
from tqdm import tqdm
from collections import deque
import ctypes
import logging_util
from deep_voice_model import PopartDeepVoice
import deep_voice_data
import conf_utils
import text_utils
# set up logging
logger = logging_util.get_basic_logger('DEEP_VOICE')
so_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),
"custom_ops.so")
if os.path.exists(so_path):
ctypes.cdll.LoadLibrary(so_path)
else:
logger.warn("Could not find custom_ops.so. To enable gradient-clipping, "
"execute `make all` before running this script.")
def _get_popart_type(np_type):
return {np.float16: 'FLOAT16', np.float32: 'FLOAT'}[np_type]
def create_inputs_for_training(builder, conf):
import logging_util
import conf_utils
import custom_op_utils
import checkpoint_utils
import transducer_blocks
import transducer_builder
from transducer_optimizer import TransducerOptimizerFactory
import ema_utils
import device
from feat_proc_cpp_async import AsyncDataProcessor
import transducer_validation
from transducer_decoder import TransducerGreedyDecoder
import test_transducer
# set up logging
logger = logging_util.get_basic_logger('TRANSDUCER_TRAIN')
def _get_popart_type(np_type):
return {
np.float16: 'FLOAT16',
np.float32: 'FLOAT',
np.int32: 'INT32'
}[np_type]
def create_inputs_for_training(builder, model_conf, conf):
""" defines the input tensors for the Transformer Transducer model """
inputs = dict()
# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
import numpy as np
from numpy.ctypeslib import ndpointer
import ctypes
import time
import logging_util
from custom_op_utils import load_custom_lib
logger = logging_util.get_basic_logger('feat_proc_cpp')
class FeatProc:
""" Wrapper class for C++ feature processing functions """
def __init__(self, conf):
libc = load_custom_lib("feat_proc")
c_contig_flag = "C_CONTIGUOUS"
self.fun_featsOutSize = libc.featsOutSize
self.fun_featsOutSize.restype = ctypes.c_uint32
self.fun_featsOutSize.argtypes = [
ndpointer(ctypes.c_uint32, flags=c_contig_flag), # featsInShape
ctypes.c_uint32, # stacking
ctypes.c_uint32
] # maxSeqLen
self.fun_featProcess = libc.featProcess
self.fun_featProcess.restype = None
self.fun_featProcess.argtypes = [
ndpointer(ctypes.c_float, flags=c_contig_flag), # featsIn
ndpointer(ctypes.c_uint32, flags=c_contig_flag), # featsInShape
# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
import argparse
import numpy as np
import time
import torch
import pytest
from common.data import features
import logging_util
logger = logging_util.get_basic_logger('test')
CPP = "cpp"
CPP_TEST = "cpp_test"
CPP_ASYNC = "cpp_async"
def add_conf_args():
""" define the argument parser object """
parser = argparse.ArgumentParser()
parser.add_argument('--batch_size', type=int, default=4, help="Batch Size")
parser.add_argument('--steps', type=int, default=5, help="Number of steps")
parser.add_argument('--task',
default="cpp_async",
choices=[CPP, CPP_TEST, CPP_ASYNC],
help="Task")
return parser
def get_train_feat_proc(conf):
import popart
import time
import os
import sys
import json
from tqdm import tqdm
from collections import deque
import ctypes
import logging_util
import conf_utils
import librispeech_data
import conformer_builder
# set up logging
logger = logging_util.get_basic_logger('CONFORMER_TRAIN')
def _get_popart_type(np_type):
return {
np.float16: 'FLOAT16',
np.float32: 'FLOAT'
}[np_type]
def load_ctc_loss_lib():
# Change cwd to the CTC directory to allow us to resolve the codelet
ctc_wd = os.path.join(os.path.dirname(__file__),
"custom_operators/ctc_loss/")
os.chdir(ctc_wd)
so_path = "build/ctc_loss.so"
# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
import popart
import onnx
import logging_util
# set up logging
logger = logging_util.get_basic_logger('EMA_UTILS')
network_names = [
"transcription_network", "prediction_network", "joint_network"
]
EMA_PREFIX = "exp_mov_avg_"
def create_exp_mov_avg_weights(builder, model_weight_names, ema_factor):
""" creates variable tensors for exponential moving averages of model weights """
ema_weight_names = dict()
for nname in network_names:
ema_weight_names[nname] = []
for uname, weight_tensor in model_weight_names[nname]:
# uname is a user-given name
# weight_tensor is the tensor name given by PopART
# The ExpMovAvg op will create new weights tensors with ID "exp_mov_avg_" + weight_tensor
# These new weight tensors will hold the data for exponential moving averages
exp_avg_weight_tensor = builder.customOp(
opName="ExpMovAvg",
opVersion=1,
# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
import numpy as np
from numpy.ctypeslib import ndpointer
import ctypes
import time
from collections import deque
import logging_util
from custom_op_utils import load_custom_lib
logger = logging_util.get_basic_logger("feat_proc_cpp_async")
class FeatProcAsync:
""" Wrapper class for calling C++ feature processing functions asynchronously """
def __init__(self, conf):
libc = load_custom_lib("feat_proc")
c_contig_flag = "C_CONTIGUOUS"
self.fun_featsOutSize = libc.featsOutSize
self.fun_featsOutSize.restype = ctypes.c_uint32
self.fun_featsOutSize.argtypes = [
ndpointer(ctypes.c_uint32, flags=c_contig_flag), # featsInShape
ctypes.c_uint32, # stacking
ctypes.c_uint32
] # maxSeqLen
self.fun_init = libc.featProcInit
self.fun_init.restype = None
self.fun_init.argtypes = [
ctypes.c_bool, # doSpecaugm
# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
import os
import time
import numpy as np
import popart
import ema_utils
import logging_util
# set up logging
logger = logging_util.get_basic_logger('CHECKPOINT_UTILS')
def remove_file_if_exists(path):
if os.path.isfile(path):
os.remove(path)
def get_training_ckpt_path(checkpt_dir):
training_ckpt_path = os.path.join(checkpt_dir,
'rnnt_checkpoint_training.onnx')
return training_ckpt_path
def get_validation_ckpt_path(checkpt_dir):
validation_ckpt_path = os.path.join(checkpt_dir,
'rnnt_checkpoint_validation.onnx')
return validation_ckpt_path
from common.data import features
import common.helpers as helpers
import common.metrics as metrics
import transducer_builder
import conf_utils
import logging_util
from rnnt_reference import config
from rnnt_reference.model import RNNT
from transducer_decoder import TransducerGreedyDecoder
import device as device_module
# set up logging
logger = logging_util.get_basic_logger('TRANSDUCER_VALIDATION')
np.set_printoptions(threshold=128)
np.set_printoptions(linewidth=1024)
def _get_popart_type(np_type):
return {np.float16: 'FLOAT16', np.float32: 'FLOAT'}[np_type]
def create_inputs_for_inference(builder, model_conf, conf):
""" defines the input tensors of the transcription network for inference """
inputs = dict()
in_feats = model_conf["rnnt"][
# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
import popart
import sys
import numpy as np
import logging_util
# set up logging
logger = logging_util.get_basic_logger('TRANSDUCER_OPTIMIZER')
class TransducerOptimizerFactory:
def __init__(self, optimizer_type, base_lr, min_lr, exp_gamma, steps_per_epoch, warmup_epochs, hold_epochs,
beta1=None, beta2=None, weight_decay=None, opt_eps=None, loss_scaling=None, gradient_clipping_norm=None, max_weight_norm=None):
""" Class for creating and updating popart optimizers
:param str optimizer_type: optimizer type - 'SGD' or 'LAMB'
:param float base_lr: base learning rate
:param float min_lr: minimum learning rate
:param float exp_gamma: gamma factor for exponential lr scheduler
:param int steps_per_epoch: training steps per one epoch
:param int warmup_epochs: initial number of epochs of increasing learning rate
:param int hold_epochs: number of epochs of constant learning rate after warmup
"""
self.optimizer_type = optimizer_type
self.base_lr = base_lr
self.min_lr = min_lr
self.exp_gamma = exp_gamma
self.steps_per_epoch = steps_per_epoch
self.warmup_epochs = warmup_epochs
self.hold_epochs = hold_epochs
self.current_lr = None
# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
import numpy as np
import popart
import os
import torch
import logging_util
import conf_utils
import text_utils
import librispeech_data
import conformer_builder
import ctcdecode
# set up logging
logger = logging_util.get_basic_logger('CONFORMER_INFERENCE')
def _get_popart_type(np_type):
return {np.float16: 'FLOAT16', np.float32: 'FLOAT'}[np_type]
def create_inputs_for_inference(builder, conf):
""" defines the input tensors for the conformer model """
inputs = dict()
inputs["mel_spec_input"] = builder.addInputTensor(
popart.TensorInfo(_get_popart_type(conf.precision), [
conf.samples_per_device, conf.mel_bands,
conf.max_spectrogram_length
]), "mel_spec_input")
# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
import torch
import torch.nn.functional as F
import logging_util
from rnnt_reference.model import label_collate
# set up logging
logger = logging_util.get_basic_logger('TransducerGreedyDecoder')
class TransducerGreedyDecoder:
"""A greedy transducer decoder.
Args:
blank_idx: inded of blank symbol
max_symbols_per_step: The maximum number of symbols that can be added
to a sequence in a single time step; if set to None then there is
no limit.
"""
def __init__(self, blank_idx, max_symbols_per_step=30, max_symbol_per_sample=None, shift_labels_by_one=True):
self.blank_idx = blank_idx
assert max_symbols_per_step is None or max_symbols_per_step > 0
self.max_symbols = max_symbols_per_step
assert max_symbol_per_sample is None or max_symbol_per_sample > 0
self.max_symbol_per_sample = max_symbol_per_sample
self._SOS = -1 # start of sequence
self.shift_labels_by_one = shift_labels_by_one
