Is it possible to substitue flash Datamodule for lightning Datamodule ?

[richarddwang]

I found Prerocess provides collate, sample/batch transform that are missed in lightning Datamodule.
Can I use an flash Datamodule (thus I can use Preprocess) to work with pl.Trainer and pl.lightningmodule ?
Is there a special way to do it or they are just interchangeable ?

[edgarriba]

We haven’t experimented with that yet. But any feedback is welcomed as I believe as midterm plan is to align both api for data modules. /cc @ethanwharris

[justusschock]

Currently this is not possible, since the LightningModule does not have the necessary hooks (they are currently living within the flash.Task). However, the flash.Task is nothing else then a lightningmodule with some additional hooks and defaults, so I recommend to use this instead.

[richarddwang]

Thanks for the comment ! @justusschock

I am also considering trying flash.Task
But is the following possible in flash.Task ? (Default ml practice that don't do weight decay on parameters of norms and bias)

def configure_optimizers(...):

    param_groups = [
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if not any(nd in n for nd in no_decay)
            ],
            "weight_decay": weight_decay,
        },
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if any(nd in n for nd in no_decay)
            ],
            "weight_decay": 0.0,
        },
    ]
   optimizer = AdamW(param_groups,...)

Or you would suggest sticking with pl instead ?

[justusschock]

Should be possible. In general I think everything that's possible in pl also works in flash, only the other direction cannot be guaranteed since flash is an extension upon lightning

[richarddwang]

Got it, thanks !