View the runnable example on GitHub

Use Channels Last Memory Format in PyTorch Lightning Training#

bigdl.nano.pytorch.Trainer API extends PyTorch Lightning Trainer with multiple integrated optimizations. You could instantiate a BigDL-Nano Trainer with channels_last=True to use the channels last memory format, i.e. NHWC (batch size, height, width, channels), as an alternative way to store tensors in classic/contiguous NCHW order.

📝 Note

Before starting your PyTorch Lightning application, it is highly recommended to run source bigdl-nano-init to set several environment variables based on your current hardware. Empirically, these variables will bring big performance increase for most PyTorch Lightning applications on training workloads.

Let’s take a self-defined LightningModule (based on a ResNet-18 model pretrained on ImageNet dataset) and dataloaders to finetune the model on OxfordIIITPet dataset as an example:

[ ]:

model = MyLightningModule()
train_loader, val_loader = create_dataloaders()

      The definition of MyLightningModule and create_dataloaders can be found in the runnable example.

To use channels last memory format (NHWC tensors) for your PyTorch Lightning application, you could simply import BigDL-Nano Trainer, and set channels_last to be True:

[ ]:

from bigdl.nano.pytorch import Trainer

trainer = Trainer(max_epochs=5, channels_last=True)

📝 Note

Channels last memory format (NHWC) is currently only implemented as an alternative to 4-dimensional NCHW tensors.

To use NHWC tensors, there is no need to modify your LightningModule and dataloaders, the only change to make is setting channels_last=True.

You could then do the normal training (and evaluation) steps with NHWC tensors:

[ ]:

trainer.fit(model, train_dataloaders=train_loader)
trainer.validate(model, dataloaders=val_loader)

📚 Related Readings