Named Tensors PyTorch
Updated: Aug 2
Named Tensors in PyTorch is aimed to make the lives of developers easier. It is a promising experimental feature in PyTorch 1.4 . It allows us to name the dimensions in a tensor. This removes the pain of remembering dimensions by positions. It provides extra safety.
Let us a create a random tensor in PyTorch -
import torch # create a tensor without using names tensor_example = torch.randn((1,2,2,1)) print(tensor_example) print(tensor_example.shape)
The output will be as follows -
tensor([[[[-0.6251], [-1.5134]], [[-2.1946], [-0.4108]]]]) torch.Size([1, 2, 2, 1])
Now if we give this code to someone who is new to our code base then it would become difficult for him to understand the shape of tensor. Like which one is height, which one is width etc.
To solve this problem, we can use named tensors -
imgs = torch.randn(1, 2, 2, 1 , names=('N', 'C', 'H', 'W')) print(imgs)
Now we can access the dimension by NCHW( Batch size, Channel, Height, Width).
tensor([[[[ 1.5122], [-1.2531]], [ 1.9728], [-0.6527]]]], names=('N', 'C', 'H', 'W'))
Note - This feature in pyTorch is experimental so don't use this in production code.