PyTorch: How to get the shape of a Tensor as a list of int

Question

In numpy, V.shape gives a tuple of ints of dimensions of V.

In tensorflow V.get_shape().as_list() gives a list of integers of the dimensions of V.

In pytorch, V.size() gives a size object, but how do I convert it to ints?

Answer 1

For PyTorch v1.0 and possibly above:

>>> import torch >>> var = torch.tensor([[1,0], [0,1]])  # Using .size function, returns a torch.Size object. >>> var.size() torch.Size([2, 2]) >>> type(var.size()) <class 'torch.Size'>  # Similarly, using .shape >>> var.shape torch.Size([2, 2]) >>> type(var.shape) <class 'torch.Size'> 

You can cast any torch.Size object to a native Python list:

>>> list(var.size()) [2, 2] >>> type(list(var.size())) <class 'list'> 

---

In PyTorch v0.3 and 0.4:

Simply list(var.size()), e.g.:

>>> import torch >>> from torch.autograd import Variable >>> from torch import IntTensor >>> var = Variable(IntTensor([[1,0],[0,1]]))  >>> var Variable containing:  1  0  0  1 [torch.IntTensor of size 2x2]  >>> var.size() torch.Size([2, 2])  >>> list(var.size()) [2, 2]