Hook for processing input data in LightningCLI

[machur]

Hi, I have a model that accepts a couple of files as input and I want to run prediction on it. Since I know the file structure, I'm not passing the paths directly via Argparse for now, I just provide directory and get all required paths programatically. I want to integrate my current approach with LightningCLI, but I don't see a place to somehow parse/process input arguments.

LightningCLI implementation:

from pytorch_lightning.cli import LightningCLI, LightningArgumentParser

class FeatureExtractorCLI(LightningCLI):

    def add_arguments_to_parser(self, parser: LightningArgumentParser):
        parser.add_argument("--input_dir", help="Input directory", required=True)

        # TODO transform input_dir into 3 new file parameters that will be used in data module

    def collect_files_from_input_dir(input_dir):
        file1 = None
        file2 = None
        file3 = None
        # some code scanning input_dir to get particular file paths...
        return file1, file2, file3

Data module that operates on separate files:

class DataModule(pl.LightningDataModule):

    def __init__(self, file1, file2, file3):
        super().__init__()
        self._file1 = file1
        self._file2 = file2
        self._file3 = file3

Is there a way to combine these two approaches or I should not mixed them? I tried overriding "parse_arguments" in LightningCLI, but unfortunately "args" are null at the time. Is there a better place to hook?

[mauvilsa]

My recommendation would be that you think of LightningCLI as a class that automatically exposes signature parameters in a CLI. However, the logic to run in most cases is not implemented as part of the CLI class. The parameters come from other classes that you can even use without the CLI.

If you want to load data from a directory, then implement a data module that receives as parameter that directory. This could be a subclass of your module that gets files. Something like:

class FilesData(pl.LightningDataModule):
    def __init__(self, file1, file2, file3):
        super().__init__()
        self._file1 = file1
        self._file2 = file2
        self._file3 = file3

class DirectoryData(FilesData):
    def __init__(self, input_dir: os.PathLike):
        file1, file2, file3 = self.get_files_from_dir(input_dir)
        super().__init__(file1, file2, file3)

Then to LightningCLI you can provide datamodule_class=FilesData and subclass_mode_data=True so that both FilesData and DirectoryData can be used from the command line.

[machur]

@mauvilsa thanks for a quick answer, that's really helpful! I implemented this in a similar way, but I wanted make sure that I didn't break any PL concepts/conventions.