Sling Academy
Home/PyTorch/PyTorch tensor shape, rank, and element count

PyTorch tensor shape, rank, and element count

Last updated: April 18, 2023

The shape of a PyTorch tensor

The shape of a PyTorch tensor is the number of elements in each dimension. You can use the shape attribute or the size() method to get the shape of a tensor as a torch.Size object, which is a subclass of tuple. You can also pass an optional argument dim to the size() method to know the size of a specific dimension as an int.

Example:

import torch

t = torch.rand(9, 7, 5)

print(f"The shape of the tensor is: {t.shape}")
print(f"The size of the tensor is: {t.size()}")

# print the size of the second dimension
print(f"The size of the second dimension is: {t.size(1)}")

Output:

The shape of the tensor is: torch.Size([9, 7, 5])
The size of the tensor is: torch.Size([9, 7, 5])
The size of the second dimension is: 7

In the context of the example above, the terms “shape” and “size” are the same.

The rank (ndims) of a PyTorch tensor

The rank of a PyTorch tensor is the number of dimensions present within the tensor. It is equal to the number of indices required to uniquely select each element of the tensor.  Rank is also known as “order”, “degree”, or “ndims.”

You can use the ndim attribute or the dim() method to get the rank of a given tensor as shown in the example below:

import torch

t = torch.rand(9, 7, 5)

print(t.dim()) # 3
print(t.ndim) # 3

Counting elements in a PyTorch tensor

You can use the numel() method to count all elements in a PyTorch tensor. It returns the total number of elements in the input tensor as an int. The nelement() method is an alias for the numel() method, and will give you the same result.

Example:

import torch

t = torch.rand(9, 7, 5)

print(t.nelement())
print(t.numel())

Output:

315
315

Next Article: PyTorch: Determine the memory usage of a tensor (in bytes)

Previous Article: Ways to Create Random Tensors in PyTorch

Series: Working with Tensors in PyTorch

PyTorch

You May Also Like

  • Addressing "UserWarning: floor_divide is deprecated, and will be removed in a future version" in PyTorch Tensor Arithmetic
  • In-Depth: Convolutional Neural Networks (CNNs) for PyTorch Image Classification
  • Implementing Ensemble Classification Methods with PyTorch
  • Using Quantization-Aware Training in PyTorch to Achieve Efficient Deployment
  • Accelerating Cloud Deployments by Exporting PyTorch Models to ONNX
  • Automated Model Compression in PyTorch with Distiller Framework
  • Transforming PyTorch Models into Edge-Optimized Formats using TVM
  • Deploying PyTorch Models to AWS Lambda for Serverless Inference
  • Scaling Up Production Systems with PyTorch Distributed Model Serving
  • Applying Structured Pruning Techniques in PyTorch to Shrink Overparameterized Models
  • Integrating PyTorch with TensorRT for High-Performance Model Serving
  • Leveraging Neural Architecture Search and PyTorch for Compact Model Design
  • Building End-to-End Model Deployment Pipelines with PyTorch and Docker
  • Implementing Mixed Precision Training in PyTorch to Reduce Memory Footprint
  • Converting PyTorch Models to TorchScript for Production Environments
  • Deploying PyTorch Models to iOS and Android for Real-Time Applications
  • Combining Pruning and Quantization in PyTorch for Extreme Model Compression
  • Using PyTorch’s Dynamic Quantization to Speed Up Transformer Inference
  • Applying Post-Training Quantization in PyTorch for Edge Device Efficiency