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 settingchannels_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