def main(
nemo_file,
enemo_file,
onnx_file,
model_type="asr",
):
if model_type == "asr":
logging.info("Preparing ASR model")
model = EncDecCTCModel.restore_from(nemo_file)
elif model_type == "speech_label":
logging.info("Preparing Speech Label Classification model")
model = EncDecClassificationModel.restore_from(nemo_file)
elif model_type == "speaker":
logging.info("Preparing Speaker Recognition model")
model = EncDecSpeakerLabelModel.restore_from(nemo_file)
else:
raise NameError(
"Available model names are asr, speech_label and speaker")
logging.info("Writing onnx file")
model.export(onnx_file, onnx_opset_version=12)
logging.info("succesfully ported onnx file")
with tarfile.open(nemo_file, "r") as archive:
archive.extract("./model_config.yaml")
with tarfile.open(enemo_file, "w") as enemo_archive:
enemo_archive.add("./model_config.yaml")
copyfile(onnx_file, "model_graph.onnx")
enemo_archive.add("model_graph.onnx")
os.remove("model_graph.onnx") # cleanup extra file
def main(
nemo_file,
enemo_file,
onnx_file,
model_type='asr',
):
if model_type == 'asr':
logging.info("Preparing ASR model")
model = EncDecCTCModel.restore_from(nemo_file)
elif model_type == 'speech_label':
logging.info("Preparing Speech Label Classification model")
model = EncDecClassificationModel.restore_from(nemo_file)
elif model_type == 'speaker':
logging.info("Preparing Speaker Recognition model")
model = EncDecSpeakerLabelModel.restore_from(nemo_file)
else:
raise NameError(
"Available model names are asr, speech_label and speaker")
logging.info("Writing onnx file")
model.export(onnx_file, onnx_opset_version=12)
logging.info("succesfully ported onnx file")
with tarfile.open(nemo_file, 'r') as archive:
archive.extract('./model_config.yaml')
with tarfile.open(enemo_file, 'w') as enemo_archive:
enemo_archive.add('./model_config.yaml')
enemo_archive.addfile(tarfile.TarInfo("model_graph.onnx"),
open(onnx_file))
def main():
parser = ArgumentParser()
parser.add_argument(
"--pretrained_model",
type=str,
default="speakerrecognition_speakernet",
required=False,
help="Pass your trained .nemo model",
)
parser.add_argument(
"--finetune_config_file",
type=str,
required=True,
help="path to speakernet config yaml file to load train, validation dataset and also for trainer parameters",
)
parser.add_argument(
"--freeze_encoder",
type=bool,
required=False,
default=True,
help="True if speakernet encoder paramteres needs to be frozen while finetuning",
)
args = parser.parse_args()
if args.pretrained_model.endswith('.nemo'):
logging.info(f"Using local speaker model from {args.pretrained_model}")
speaker_model = EncDecSpeakerLabelModel.restore_from(restore_path=args.pretrained_model)
elif args.pretrained_model.endswith('.ckpt'):
logging.info(f"Using local speaker model from checkpoint {args.pretrained_model}")
speaker_model = EncDecSpeakerLabelModel.load_from_checkpoint(checkpoint_path=args.pretrained_model)
else:
logging.info("Using pretrained speaker recognition model from NGC")
speaker_model = EncDecSpeakerLabelModel.from_pretrained(model_name=args.pretrained_model)
finetune_config = OmegaConf.load(args.finetune_config_file)
if 'test_ds' in finetune_config.model and finetune_config.model.test_ds is not None:
finetune_config.model.test_ds = None
logging.warning("Removing test ds")
speaker_model.setup_finetune_model(finetune_config.model)
finetune_trainer = pl.Trainer(**finetune_config.trainer)
speaker_model.set_trainer(finetune_trainer)
_ = exp_manager(finetune_trainer, finetune_config.get('exp_manager', None))
speaker_model.setup_optimization(finetune_config.optim)
if args.freeze_encoder:
for param in speaker_model.encoder.parameters():
param.requires_grad = False
finetune_trainer.fit(speaker_model)
def main(
nemo_file, onnx_file, model_type='asr',
):
if model_type == 'asr':
logging.info("Preparing ASR model")
model = EncDecCTCModel.restore_from(nemo_file)
elif model_type == 'speech_label':
logging.info("Preparing Speech Label Classification model")
model = EncDecClassificationModel.restore_from(nemo_file)
elif model_type == 'speaker':
logging.info("Preparing Speaker Recognition model")
model = EncDecSpeakerLabelModel.restore_from(nemo_file)
else:
raise NameError("Available model names are asr, speech_label and speaker")
logging.info("Writing onnx file")
model.export(onnx_file, onnx_opset_version=12)
logging.info("succesfully ported onnx file")
def main(cfg):
logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
device = 'cuda' if torch.cuda.is_available() else 'cpu'
enrollment_manifest = cfg.data.enrollment_manifest
test_manifest = cfg.data.test_manifest
out_manifest = cfg.data.out_manifest
sample_rate = cfg.data.sample_rate
backend = cfg.backend.backend_model.lower()
if backend == 'cosine_similarity':
model_path = cfg.backend.cosine_similarity.model_path
batch_size = cfg.backend.cosine_similarity.batch_size
if model_path.endswith('.nemo'):
speaker_model = EncDecSpeakerLabelModel.restore_from(model_path)
else:
speaker_model = EncDecSpeakerLabelModel.from_pretrained(model_path)
enroll_embs, _, enroll_truelabels, enroll_id2label = EncDecSpeakerLabelModel.get_batch_embeddings(
speaker_model,
enrollment_manifest,
batch_size,
sample_rate,
device=device,
)
test_embs, _, _, _ = EncDecSpeakerLabelModel.get_batch_embeddings(
speaker_model,
test_manifest,
batch_size,
sample_rate,
device=device,
)
# length normalize
enroll_embs = enroll_embs / (np.linalg.norm(
enroll_embs, ord=2, axis=-1, keepdims=True))
test_embs = test_embs / (np.linalg.norm(
test_embs, ord=2, axis=-1, keepdims=True))
# reference embedding
reference_embs = []
keyslist = list(enroll_id2label.keys())
for label_id in keyslist:
indices = np.where(enroll_truelabels == label_id)
embedding = (enroll_embs[indices].sum(
axis=0).squeeze()) / len(indices)
reference_embs.append(embedding)
reference_embs = np.asarray(reference_embs)
scores = np.matmul(test_embs, reference_embs.T)
matched_labels = scores.argmax(axis=-1)
elif backend == 'neural_classifier':
model_path = cfg.backend.neural_classifier.model_path
batch_size = cfg.backend.neural_classifier.batch_size
if model_path.endswith('.nemo'):
speaker_model = EncDecSpeakerLabelModel.restore_from(model_path)
else:
speaker_model = EncDecSpeakerLabelModel.from_pretrained(model_path)
featurizer = WaveformFeaturizer(sample_rate=sample_rate)
dataset = AudioToSpeechLabelDataset(
manifest_filepath=enrollment_manifest,
labels=None,
featurizer=featurizer)
enroll_id2label = dataset.id2label
if speaker_model.decoder.final.out_features != len(enroll_id2label):
raise ValueError(
"number of labels mis match. Make sure you trained or finetuned neural classifier with labels from enrollement manifest_filepath"
)
_, test_logits, _, _ = EncDecSpeakerLabelModel.get_batch_embeddings(
speaker_model,
test_manifest,
batch_size,
sample_rate,
device=device,
)
matched_labels = test_logits.argmax(axis=-1)
with open(test_manifest, 'rb') as f1, open(out_manifest,
'w',
encoding='utf-8') as f2:
lines = f1.readlines()
for idx, line in enumerate(lines):
line = line.strip()
item = json.loads(line)
item['infer'] = enroll_id2label[matched_labels[idx]]
json.dump(item, f2)
f2.write('\n')
logging.info(
"Inference labels have been written to {} manifest file".format(
out_manifest))
def main():
parser = ArgumentParser()
parser.add_argument(
"--spkr_model", type=str, default="titanet_large", required=True, help="Pass your trained .nemo model",
)
parser.add_argument(
"--train_manifest", type=str, required=True, help="path to train manifest file to match labels"
)
parser.add_argument(
"--test_manifest", type=str, required=True, help="path to test manifest file to perform inference"
)
parser.add_argument("--batch_size", type=int, default=32)
args = parser.parse_args()
torch.set_grad_enabled(False)
if args.spkr_model.endswith('.nemo'):
logging.info(f"Using local speaker model from {args.spkr_model}")
speaker_model = EncDecSpeakerLabelModel.restore_from(restore_path=args.spkr_model)
else:
logging.error(f"Please pass a trained .nemo file")
sys.exit()
labels = []
with open(args.train_manifest, 'rb') as f:
lines = f.readlines()
for line in lines:
line = line.strip()
item = json.loads(line)
labels.append(item['label'])
labels_map = sorted(set(labels))
label2id, id2label = {}, {}
for label_id, label in enumerate(labels_map):
label2id[label] = label_id
id2label[label_id] = label
speaker_model.setup_test_data(
test_data_layer_params={
'sample_rate': 16000,
'manifest_filepath': args.test_manifest,
'labels': labels_map,
'batch_size': args.batch_size,
'trim_silence': False,
'shuffle': False,
}
)
if can_gpu:
speaker_model = speaker_model.cuda()
speaker_model.eval()
speaker_model.test_dataloader()
all_labels = []
all_logits = []
for test_batch in tqdm(speaker_model.test_dataloader()):
if can_gpu:
test_batch = [x.cuda() for x in test_batch]
with autocast():
audio_signal, audio_signal_len, labels, _ = test_batch
logits, _ = speaker_model.forward(input_signal=audio_signal, input_signal_length=audio_signal_len)
all_logits.extend(logits.cpu().numpy())
all_labels.extend(labels.cpu().numpy())
all_logits, true_labels = np.asarray(all_logits), np.asarray(all_labels)
infer_labels = all_logits.argmax(axis=1)
out_manifest = os.path.basename(args.test_manifest).split('.')[0] + '_infer.json'
out_manifest = os.path.join(os.path.dirname(args.test_manifest), out_manifest)
with open(args.test_manifest, 'rb') as f1, open(out_manifest, 'w', encoding='utf-8') as f2:
lines = f1.readlines()
for idx, line in enumerate(lines):
line = line.strip()
item = json.loads(line)
item['infer'] = id2label[infer_labels[idx]]
json.dump(item, f2)
f2.write('\n')
logging.info("Inference labels have been written to {} manifest file".format(out_manifest))