How to change pytorch DataLoader/Dataset for nemo?

[EndOfTheGlory]

Hello. I want to create custom DataLoader or Dataset that will change text of data into different tokens for Marblenet model on the fly (i.e. while training). For that purpose I researched how DataLoader and Dataset work both in pytorch and pytorch lightning. But maybe because it still hasn't settled in, I don't understand how to do everything using nemo library models, considering that for nemo models to work we use config files with to create dataloaders automatically and I can't understand what we can do with them to work correctly with nemo models.

So, for my point to be clear perhaps I should be a little verbose.

The only threads I found about this discussion are following: this and this. But both of them do little to make me understand what should be done.

I researched the following:

As pytorch lightning says in succinct tutorial that it uses natives classes of pytorch for Dataloaders. For vanilla pytorch there is another article how to create custom Dataloaders and Datasets. For pytorch lightning I also found class DataModules that we could use for using custom dataloaders (here).

But here comes the problem. Let's look at my code draft for reference:

clearml_logger = task.get_logger()
val_samples = model.val_dataloader()
trainer = pl.Trainer(devices='auto',
                accelerator='auto',
                strategy='ddp',
                max_epochs=args.max_epochs,
                precision=args.precision,
                callbacks = [TextRecognitionLogger(val_samples, clearml_logger)],
                accumulate_grad_batches=args.accumulate_grad_batches,
                enable_checkpointing=False,
                logger=False,
                enable_progress_bar=True,
                val_check_interval=args.val_check_interval,
                fast_dev_run=2)
model.set_trainer(trainer)

experiment_cfg = exp_manager.ExpManagerConfig(
    exp_dir=f'experiments/adapters/',
    name=f"ASR-Adapter",
    resume_if_exists=True,
    resume_ignore_no_checkpoint=True,
    checkpoint_callback_params=exp_manager.CallbackParams(
        monitor="val_wer",
        mode="min",
        save_top_k=1,
        always_save_nemo=True,
        save_best_model=True
        )
    )
experiment_cfg = OmegaConf.structured(experiment_cfg)
logdir = exp_manager.exp_manager(trainer, experiment_cfg)
trainer.fit(model)

In the code model is initialized as following:

config_path = '/usr/src/app/configs/marblenet_3x2x64_ru_vocab.yaml'
cfg = OmegaConf.load(config_path)

cfg["model"]["train_ds"]["manifest_filepath"] = TRAIN_DS
cfg["model"]["validation_ds"]["manifest_filepath"] = VAL_DS
cfg["model"]["test_ds"]["manifest_filepath"] = TEST_DS
model = nemo_asr.models.EncDecCTCModel(cfg=cfg["model"])
cfg = copy.deepcopy(model.cfg)

# some changes done to config ...

So, the model is created using config in the style nemo requires, and through code I give paths to datasets that are used in training/testing/validation. We can grab dataloaders using nemo's val_samples = model.val_dataloader(), but... what to do next?

I mean, I could probably change dataloader creating my_custom_dataloader using the paths to data like I give to nemo and than do model.val_dataloader() = my_custom_dataloader, but to me it seems very crude, inelegant and plain wrong decision. Because as was said in first issue by more knowledgeable user:
The data loaders are created automatically when you create a model by passing the config as it contains the data loaders' info. No need to create them separately.

So it's seems as wrong decision to override model properties.

I tried looking in pytorch lightning Trainer class documentation, but it seems like a wrong move to hope on trainer also, as was advised by another user. The only parameter I found that could be used to change dataloader is use_distributed_sampler using native pytorch class.

The only option left was to pass dataloaders as arguments to trainer.fit(here's documentation), but first problem is that we can only pass two dataloaders: train_dataloaders and val_dataloaders, what to do with test_dataloaders? And finally, how will nemo behave, considering we use nemo model for fitting the trainer and in reverse?

(look full snippet higher)

model.set_trainer(trainer)

# exp_manager.ExpManagerConfig actions ...

trainer.fit(model)

What should I do to make it all work?

I hope you can help me settle my confusion. If something is unclear, ask me any question and I will try to my best abilities to answer.

[EndOfTheGlory]

Well, for anyone looking at this in the future, I have two advices how to solve the problem.

First one is depicted in one NeMo tutorial (https://github.com/NVIDIA/NeMo/blob/main/tutorials/01_NeMo_Models.ipynb), where they create custom Dataset and model class.

Second one is to create custom model class by redefining it's parent classes like we have done with colleague. It means that you should find the function that processes what you need and redifine it to proccess it your way. In my case it was class ASRAudioText (I changed how it processes text via my_text = re.sub('[a-zA-Zа-яА-Я]', '-', item['text']) for different tokens).

class ASRAudioText(collections.AudioText):
    """`AudioText` collector from asr structured json files."""

    def __init__(self, manifests_files: Union[str, List[str]], *args, **kwargs):
        """Parse lists of audio files, durations and transcripts texts.

        Args:
            manifests_files: Either single string file or list of such -
                manifests to yield items from.
            *args: Args to pass to `AudioText` constructor.
            **kwargs: Kwargs to pass to `AudioText` constructor.
        """
        ids, audio_files, durations, texts, offsets, = (
            [],
            [],
            [],
            [],
            [],
        )
        speakers, orig_srs, token_labels, langs = [], [], [], []
        for item in manifest.item_iter(manifests_files):
            ids.append(item['id'])
            audio_files.append(item['audio_file'])
            
            durations.append(item['duration'])
            my_text = re.sub('[a-zA-Zа-яА-Я]', '-', item['text'])
            
            texts.append(my_text)
            offsets.append(item['offset'])
            speakers.append(item['speaker'])
            orig_srs.append(item['orig_sr'])
            token_labels.append(item['token_labels'])
            langs.append(item['lang'])
        super().__init__(
            ids, audio_files, durations, texts, offsets, speakers, orig_srs, token_labels, langs, *args, **kwargs
        )

Than we must inherit this changed class in all other function all the way till our custom class.
For reference, in my case I used my redefined classes in the following sequence: get_char_dataset -> ASRManifestProcessor -> _AudioTextDataset(Dataset) -> AudioToCharDataset(_AudioTextDataset) -> MYEncDecCTCModel(EncDecCTCModel). In this classes we use our previous defined classes instead of nemo library ones.

I hope what I said makes sense and helps some poor soul.

[eduardo-onate]

Thank you very much, you indirectly helped this poor soul.

My case was a bit different. I share it hoping it could be helpful for other fellow users.

I'm using an EncDecSpeakerLabelModel for ASR and I was getting ValueError: Expected more than 1 value per channel when training, got input size torch.Size([1, 6144, 1]). According to this helpful discussion on PyTorch's forum I had to set the flag drop_last=True while constructing the DataLoaders. Since NeMo does all of this from behind I followed the path of the error and found a file where that parameter was asigned as drop_last=config.get('drop_last', False). The solution was then as simple as adding drop_last: True in the sections of train_ds and validation_ds in the config file.

train_ds:
        ...
	drop_last: True
	...
...
validation_ds:
        ...
	drop_last: True
	...

Good luck!