How to use NEMO - 10 common examples

# parser.add_argument('-evalp', '--evalp',
#                     help='evaluation period', default=1)
# parser.add_argument('-an', '--addName',
#                     help='An add name for the save folder', default='')
# parser.add_argument('-eb', '--eval_batch', help='Evaluation Batch_size',
#                     type=int, default=0)

args = parser.parse_args()
DOMAINS = {"attraction": 0, "restaurant": 1, "taxi": 2, "train": 3, "hotel": 4}

if not os.path.exists(args.data_dir):
    raise ValueError(f'Data not found at {args.data_dir}')

work_dir = f'{args.work_dir}/{args.dataset_name.upper()}'

nf = nemo.core.NeuralModuleFactory(backend=nemo.core.Backend.PyTorch,
                                   local_rank=args.local_rank,
                                   optimization_level=args.amp_opt_level,
                                   log_dir=work_dir,
                                   create_tb_writer=True,
                                   files_to_copy=[__file__],
                                   add_time_to_log_dir=True)

data_layer = nemo_nlp.WOZDSTDataLayer(args.data_dir,
                                      DOMAINS,
                                      batch_size=args.batch_size,
                                      mode='train')
src_ids, src_lens, tgt_ids, tgt_lens, gate_labels, turn_domain = data_layer()
vocab_size = len(data_layer._dataset.vocab)
steps_per_epoch = len(data_layer) // args.batch_size

encoder = EncoderRNN(vocab_size,

print("Loading checkpoints...")
    if pre_v09_model:
        print("  Converting pre v0.9 checkpoint...")
        ckpt = torch.load(nn_encoder)
        new_ckpt = {}
        for k, v in ckpt.items():
            new_k = k.replace('.conv.', '.mconv.')
            if len(v.shape) == 3:
                new_k = new_k.replace('.weight', '.conv.weight')
            new_ckpt[new_k] = v
        jasper_encoder.load_state_dict(new_ckpt)
    else:
        jasper_encoder.restore_from(nn_encoder)
    jasper_decoder.restore_from(nn_decoder)

    nf = nemo.core.NeuralModuleFactory(create_tb_writer=False)
    print("Exporting encoder...")
    nf.deployment_export(jasper_encoder, nn_onnx_encoder,
                         nemo.core.neural_factory.DeploymentFormat.ONNX,
                         torch.zeros(batch_size,
                                     num_encoder_input_features,
                                     time_steps,
                                     dtype=torch.float, device="cuda:0"))
    print("Exporting decoder...")
    nf.deployment_export(jasper_decoder, nn_onnx_decoder,
                         nemo.core.neural_factory.DeploymentFormat.ONNX,
                         (torch.zeros(batch_size,
                                      num_decoder_input_features,
                                      time_steps // 2,
                                      dtype=torch.float, device="cuda:0")))
    print("Export completed successfully.")

opt_level)


name = construct_name('ZeroDS-Jasper10x5', lr, batch_size, num_gpus, num_epochs,
                      weight_decay)

tb_writer = SummaryWriter(name)

if args.local_rank is not None:
    device = nemo.core.DeviceType.AllGpu
    print('Doing ALL GPU')
else:
    device = nemo.core.DeviceType.GPU

# instantiate Neural Factory with supported backend
neural_factory = nemo.core.NeuralModuleFactory(
    backend=nemo.core.Backend.PyTorch,
    local_rank=args.local_rank,
    optimization_level=opt_level,
    placement=device)

jasper_model_definition = toml.load("../../examples/nemo_asr/jasper10x5.toml")
jasper_model_definition['placement'] = device
labels = jasper_model_definition['labels']['labels']

# train_manifest = "/mnt/D1/Data/librispeech/librivox-train-all.json"
train_manifest = args.train_manifest

featurizer_config = jasper_model_definition['input']
data_preprocessor = neural_factory.get_module(
        name="AudioToMelSpectrogramPreprocessor",
        collection="nemo_asr",

length=p_length_e1)
log_probs_e1 = jasper_decoder(encoder_output=encoded_e1)
predictions_e1 = greedy_decoder(log_probs=log_probs_e1)
loss_e1 = ctc_loss(log_probs=log_probs_e1,
                   targets=transcript_e1,
                   input_length=encoded_len_e1,
                   target_length=transcript_len_e1)


print('\n\n\n================================')
print("Total number of parameters: {0}".format(
    jasper_decoder.num_weights + jasper_encoder.num_weights))
print('================================')

# Callbacks needed to print info to console and Tensorboard
train_callback = nemo.core.SimpleLossLoggerCallback(
    tensor_list2str=lambda x: str(x[0].item()),
    tb_writer=tb_writer,
    tensor_list2str_evl=lambda x: monitor_asr_train_progress(x, labels=labels))

eval_callback1 = nemo.core.EvaluatorCallback(
    eval_tensors=[loss_e1, predictions_e1, transcript_e1, transcript_len_e1],
    user_iter_callback=lambda x, y: process_evaluation_batch(
        x, y, labels=labels),
    user_epochs_done_callback=lambda x: process_evaluation_epoch(x,
                                                                 tag="DEV-CLEAN"),
    eval_step=500,
    tb_writer=tb_writer)


def lr_policy(initial_lr, step, N):
    res = initial_lr * ((N - step + 1) / N) ** 2

name = construct_name('ZeroDS-Jasper10x5', lr, batch_size, num_gpus, num_epochs,
                      weight_decay)

tb_writer = SummaryWriter(name)

if args.local_rank is not None:
    device = nemo.core.DeviceType.AllGpu
    print('Doing ALL GPU')
else:
    device = nemo.core.DeviceType.GPU

# instantiate Neural Factory with supported backend
neural_factory = nemo.core.NeuralModuleFactory(
    backend=nemo.core.Backend.PyTorch,
    local_rank=args.local_rank,
    optimization_level=opt_level,
    placement=device)

jasper_model_definition = toml.load("../../examples/nemo_asr/jasper10x5.toml")
jasper_model_definition['placement'] = device
labels = jasper_model_definition['labels']['labels']

# train_manifest = "/mnt/D1/Data/librispeech/librivox-train-all.json"
train_manifest = args.train_manifest

featurizer_config = jasper_model_definition['input']
data_preprocessor = neural_factory.get_module(
        name="AudioToMelSpectrogramPreprocessor",
        collection="nemo_asr",
        params=featurizer_config)

opt_level)


name = construct_name('Jasper10x5', lr, batch_size, num_gpus, num_epochs,
                      weight_decay)
tb_writer = SummaryWriter(name)

if args.local_rank is not None:
    device = nemo.core.DeviceType.AllGpu
    print('Doing ALL GPU')
else:
    device = nemo.core.DeviceType.GPU

# instantiate Neural Factory with supported backend
neural_factory = nemo.core.NeuralModuleFactory(
    backend=nemo.core.Backend.PyTorch,
    local_rank=args.local_rank,
    optimization_level=opt_level,
    placement=device)

jasper_model_definition = toml.load("../../examples/nemo_asr/jasper10x5.toml")
jasper_model_definition['placement'] = device
labels = jasper_model_definition['labels']['labels']

train_manifest = "/mnt/D1/Data/librispeech/librivox-train-all.json"
#train_manifest = args.train_manifest
val_manifest1 = "/mnt/D1/Data/librispeech/librivox-dev-clean.json"
# val_manifest2 = "/mnt/D1/Data/librispeech/librivox-dev-other.json"
#val_manifest1 = args.val_manifest1


featurizer_config = jasper_model_definition['input']

normal_dl = nemo_asr.AudioToTextDataLayer(
            featurizer_config=self.featurizer_config,
            manifest_filepath=self.manifest_filepath,
            labels=self.labels,
            batch_size=batch_size,
            placement=DeviceType.GPU,
            drop_last=True,
            shuffle=False
        )
        trimmed_dl = nemo_asr.AudioToTextDataLayer(
            featurizer_config=self.featurizer_config,
            manifest_filepath=self.manifest_filepath,
            trim_silence=True,
            labels=self.labels,
            batch_size=batch_size,
            placement=DeviceType.GPU,
            drop_last=True,
            shuffle=False
        )
        for norm, trim in zip(normal_dl.data_iterator,
                              trimmed_dl.data_iterator):
            for point in range(batch_size):
                self.assertTrue(norm[1][point].data >= trim[1][point].data)

return "{0}-lr_{1}-bs_{2}x{3}-e_{4}-wd_{5}-OPT-{6}".format(name, lr,
                                                               batch_size,
                                                               num_gpus,
                                                               num_epochs, wd,
                                                               opt_level)


name = construct_name('Jasper10x5', lr, batch_size, num_gpus, num_epochs,
                      weight_decay)
tb_writer = SummaryWriter(name)

if args.local_rank is not None:
    device = nemo.core.DeviceType.AllGpu
    print('Doing ALL GPU')
else:
    device = nemo.core.DeviceType.GPU

# instantiate Neural Factory with supported backend
neural_factory = nemo.core.NeuralModuleFactory(
    backend=nemo.core.Backend.PyTorch,
    local_rank=args.local_rank,
    optimization_level=opt_level,
    placement=device)

jasper_model_definition = toml.load("../../examples/nemo_asr/jasper10x5.toml")
jasper_model_definition['placement'] = device
labels = jasper_model_definition['labels']['labels']

train_manifest = "/mnt/D1/Data/librispeech/librivox-train-all.json"
#train_manifest = args.train_manifest
val_manifest1 = "/mnt/D1/Data/librispeech/librivox-dev-clean.json"
# val_manifest2 = "/mnt/D1/Data/librispeech/librivox-dev-other.json"

batch_size=eval_batch_size,
            num_workers=cpu_per_dl,
            **eval_dl_params,
        )

        audio, audio_len, = data_layer_eval()
        spec_target, spec_target_len = data_preprocessor(
            input_signal=audio,
            length=audio_len)

        audio_pred, log_s_list, log_det_W_list = waveglow(
            mel_spectrogram=spec_target, audio=audio)

        # create corresponding eval callback
        tagname = os.path.basename(eval_dataset).split(".")[0]
        eval_callback = nemo.core.EvaluatorCallback(
            eval_tensors=[audio_pred, spec_target, spec_target_len],
            user_iter_callback=waveglow_process_eval_batch,
            user_epochs_done_callback=lambda x: x,
            tb_writer_func=partial(
                waveglow_eval_log_to_tb_func,
                tag=tagname,
                mel_fb=data_preprocessor.filter_banks),
            eval_step=eval_freq,
            tb_writer=neural_factory.tb_writer)

        callbacks.append(eval_callback)
    return callbacks

nf = nemo.core.NeuralModuleFactory()

# instantiate necessary neural modules
# RealFunctionDataLayer defaults to f=torch.sin, sampling from x=[-4, 4]
dl = nemo.tutorials.RealFunctionDataLayer(
    n=10000, batch_size=128)
fx = nemo.tutorials.TaylorNet(dim=4)
loss = nemo.tutorials.MSELoss()

# describe activation's flow
x, y = dl()
p = fx(x=x)
lss = loss(predictions=p, target=y)

# SimpleLossLoggerCallback will print loss values to console.
callback = nemo.core.SimpleLossLoggerCallback(
    tensors=[lss],
    print_func=lambda x: print(f'Train Loss: {str(x[0].item())}'))

# Invoke "train" action
nf.train([lss], callbacks=[callback],
         optimization_params={"num_epochs": 3, "lr": 0.0003},
         optimizer="sgd")