PyTorch Tensors

PyTorch tensors are the basic data structure of every model: inputs, weights, outputs and gradients are all tensors. A tensor is a multi-dimensional array of numbers with a fixed data type, very similar to a NumPy array but able to run on a GPU and track gradients.

Creating Tensors

import torch

scalar = torch.tensor(7)
vector = torch.tensor([1, 2, 3])
matrix = torch.tensor([[1.0, 2.0], [3.0, 4.0]])

print(scalar.ndim, vector.ndim, matrix.ndim)
print(matrix.shape)
0 1 2
torch.Size([2, 2])

ndim is the number of dimensions and shape is the size of each one. An image batch, for example, is usually 4-D: [batch, channels, height, width].

Shape, dtype and device

Every tensor has three key attributes. Mismatches between them cause most beginner errors.

t = torch.tensor([1, 2, 3])
print(t.dtype, t.device)

f = t.float()
print(f.dtype)
print(torch.tensor([1.5, 2]).dtype)
torch.int64 cpu
torch.float32
torch.float32
  • float32 is the default for decimals and the usual dtype for model weights.
  • int64 (long) is the default for integers and is used for class labels.
  • bfloat16 and float16 save memory on modern GPUs.
  • bool tensors come from comparisons like t > 2.

Factory Functions

print(torch.zeros(2, 3))
print(torch.ones(2, dtype=torch.int32))
print(torch.arange(0, 10, 3))
print(torch.linspace(0, 1, 5))
torch.manual_seed(0)
print(torch.rand(2, 2))
tensor([[0., 0., 0.],
        [0., 0., 0.]])
tensor([1, 1], dtype=torch.int32)
tensor([0, 3, 6, 9])
tensor([0.0000, 0.2500, 0.5000, 0.7500, 1.0000])
tensor([[0.4963, 0.7682],
        [0.0885, 0.1320]])

torch.manual_seed makes random numbers repeatable, which is essential when debugging. torch.randn draws from a normal distribution instead of a uniform one.

Indexing and Slicing

x = torch.arange(12).reshape(3, 4)
print(x)
print(x[0])        # first row
print(x[:, 1])     # second column
print(x[1, 2].item())
print(x[x > 8])    # boolean mask
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]])
tensor([0, 1, 2, 3])
tensor([1, 5, 9])
6
tensor([ 9, 10, 11])

.item() converts a one-element tensor to a plain Python number, which you will use constantly to log loss values.

Converting to and from NumPy

import numpy as np

a = np.array([1.0, 2.0])
t = torch.from_numpy(a)
a[0] = 99
print(t)
print(t.numpy())
tensor([99.,  2.], dtype=torch.float64)
[99.  2.]

from_numpy shares memory with the array, so changing one changes the other. Use torch.tensor(a) when you want an independent copy. Note NumPy defaults to float64; call .float() before feeding it to a model.

Frequently asked questions

What is a tensor in PyTorch?

A multi-dimensional array of numbers with a single dtype that can live on the CPU or GPU and optionally track gradients for training.

What is the difference between torch.tensor and torch.Tensor?

torch.tensor(data) infers the dtype from the data and is the recommended way to create tensors. torch.Tensor is the class itself and always produces float32, which can be surprising.

How do I convert a PyTorch tensor to a NumPy array?

Call t.numpy() for CPU tensors. For GPU tensors or tensors that require grad, use t.detach().cpu().numpy().