Looking at the tutorial for the channels last format available at (beta) Channels Last Memory Format in PyTorch 4, I decided to try it out, but found myself facing a roughly x6 performance penalty rather than any gain. Looking at it, I feel that this is because the memory reorganisation was happening on the GPU–the only time the Input variable is directly exposed is within the training loop. Obviously, restructuring the data once it’s loaded onto the GPU is undesirable.
However, does is there any way to change the channel ordering when working with the DataLoader class or using ImageLoader to load data? Neither class has its own memory format due to being iterables, and passing in lambda x: x.to(memory_format=torch.channels_last) or lambda x: x.contiguous(memory_format=torch.channels_last) to a transforms.Lambda function results in this error:
AttributeError: Can't pickle local object 'main.<locals>.<lambda>'
It would seem that a unique transform method would be needed?