How to implement channels last memory format callback

[hal-314]

Hi there

Pytorch docs recommends using channels last when training vision models in mixed precision. To enable, you need to do two changes:

1. Move you model to channel last format: model = model.to(memory_format=torch.channels_last) # Replace with your model. This can be done in on_fit_start callback hook as model.to performs an in place modification:

class ChannelsLast(pl.Callback):
    def on_fit_start(self, trainer, pl_module: pl.LightningModule) -> None:
        # Inplace model modification
        pl_module.to(memory_format=torch.channels_last)

2. Move input data to channel last format before feeding to the model: input = input.to(memory_format=torch.channels_last).

My problem is in step 2. I don't find any PyTorch lightning hook that allows me to make this modification to the batch :/. The only options left are to add it as data transforms (that must be used in conjunction with the callback) or doing all channel last related logic inside the LightningModule. I would prefer to avoid this last solution as it could clutter the LightningModule with unnecessary code.

Do you know a to do step 2 in a callback?

[carmocca]

You can use any of the following if done inside the LightningModule:

batch = self.on_before_batch_transfer(batch, dataloader_idx)
batch = self.transfer_batch_to_device(batch, device)
batch = self.on_after_batch_transfer(batch, dataloader_idx)

If you really need to do it in the Callback, I guess you could use on_train_batch_start since the modification is in-place. But I wouldn't recommend it.