Processing in predict_step() requires access to a DataModule attribute

[dianemarquette]

Hi!

In my LightningDataModule, I apply preprocessing transforms to my input data before feeding it to my dataloader. In the same datamodule, I also defined the postprocessing transforms to apply after the inference.

self.post_transforms = Compose([
            AsDiscreted(keys="pred", argmax=True, to_onehot=False, n_classes=3),
            Invertd(
                keys="pred",  # invert the `pred` data field, also support multiple fields
                transform=val_transforms,
                loader=val_dataloader,
                orig_keys="img",  # get the previously applied pre_transforms information on the `img` data field,
                                  # then invert `pred` based on this information. we can use same info
                                  # for multiple fields, also support different orig_keys for different fields
                meta_keys="pred_meta_dict",  # key field to save inverted meta data, every item maps to `keys`
                orig_meta_keys="img_meta_dict",  # get the meta data from `img_meta_dict` field when inverting,
                                                 # for example, may need the `affine` to invert `Spacingd` transform,
                                                 # multiple fields can use the same meta data to invert
                meta_key_postfix="meta_dict",  # if `meta_keys=None`, use "{keys}_{meta_key_postfix}" as the meta key,
                                               # if `orig_meta_keys=None`, use "{orig_keys}_{meta_key_postfix}",
                                               # otherwise, no need this arg during inverting
                nearest_interp=True,  # change to use "nearest" mode in interpolation when inverting
                to_tensor=True,  # convert to PyTorch Tensor after inverting
            ),
            SaveImaged(keys="pred", meta_keys="pred_meta_dict", output_dir="./out", output_postfix="seg", resample=False),
        ])

I want to apply these post_transforms to my inference outputs in predict_step(). What would be the best "Lightning way" to give access to my datamodule.post_transforms attribute to predict_step?

def predict_step(self, batch: Any, batch_idx: int):
        batch["pred"] = self.forward(batch)
        post_transforms(batch)

Thanks in advance for your suggestions :)

[awaelchli]

Hi, you should be able to access the datamodule inside the LightningModule.
Try

def predict_step(self, batch: Any, batch_idx: int):
        batch["pred"] = self(batch)
        self.datamodule.post_transforms(batch)

Also, another tipp: Better use self() instead of self.forward() (generally in PyTorch).

[dianemarquette]

Ok, thanks for your reply and the tip regarding self()! I added trainer before datamodule and now it works great:

self.trainer.datamodule.post_transforms(batch)

[abuwady]

Ok, thanks for your reply and the tip regarding self()! I added trainer before datamodule and now it works great:

self.trainer.datamodule.post_transforms(batch)

[email protected]