Almacenamiento de tensores en PyTorch

Como hemos visto, los tensores son un elemento básico en PyTorch. Son los elementos de diferente rango (escalares, vectores, matrics, etc.) que contienen la información en PyTorch y sobre los que actúan los modelos.

Sin embargo, a pesar de su rango y dimensiones, los tensores se almacenan en memoria en un array unidimensional de elementos contiguos. De esta forma se homogeiniza su almacenamiento y se hace más eficiente su tratamient

A esta información se accede a través de la clase Torch.storage(). A continuación, vemos un ejemplo de un tensor y su almacenamiento en un array unidimensional.

import torch 

tensor1 = torch.tensor([[2.0, 1.0], [4.0, 3.0], [2.0, 1.0]])
tensor1.storage()

 2.0
 1.0
 4.0
 3.0
 2.0
 1.0

Los tensores son vistas de este almacenamiento y para poder relacionar las diferentes vistas con un único almacenamiento, PyTorch hace uso de los metadatos size, offset y stride.

Size es como shape en numpy e indica el número de elementos en cada dimensión. Offset es el índice en el almacenamiento (storage) que indica donde está el primer elemento del tensor y el stride es el número de desplazamientos en el almacenamiento necesarios para indexar el siguiente elemento en cada dimensión.

A continuación, vemos estos metadatos para nuestro tensor.

print("El tamaño (size) es", tensor1.size())

print("El offset es", tensor1.storage_offset())

print("El stride es", tensor1.stride())

El tamaño (size) es torch.Size([3, 2])
El offset es 0
El stride es (2, 1)

Como se puede ver, nuestro tensor tiene un tamaño de 3 filas y dos columnas y el primer elemento del tensor está en la posición cero del almacenamiento.

Para movernos una posición en la primera dimensión/segunda dimensión (fila/columna) necesitamos avanzar (2,1) posiciones en el almacenamiento.

A continuación creamos otro tensor de 3 filas y 3 columnas y vemos como ahora para movernos en la primera dimensión (fila) necesitamos avanzar 3 posiciones en el almacenamiento.

tensor2=torch.rand(3,3)

print("El tamaño (size) es", tensor2.size())

print("El offset es", tensor2.storage_offset())

print("El stride es", tensor2.stride())

El tamaño (size) es torch.Size([3, 3])
El offset es 0
El stride es (3, 1)

Un tensor cuya vista coincida con el almacenamiento es un tensor contiguo. Es decir, sus valores se almacenan desde la dimensión más a la derecha hacia adelante. Esto es conveniente en términos de eficiencia porque podemos recorrer el tensor en orden sin hacer saltos en la memoria.

Hay una serie de operaciones en PyTorch que no cambien el almacenamiento pero sí que cambian la vista y consecuentemente los metadatos de vista del tensor. Son las operaciones narrow(), view(), expand() y transpose().

Vamos a transponer nuestro primer vector y ver su storage y sus metadatos:

tensor1_transpose=tensor1.transpose(0,1)
print(tensor1_transpose)
print(tensor1_transpose.storage())

print("¿Es el tensor contiguo?", tensor1_transpose.is_contiguous())

print("El tamaño (size) es", tensor1_transpose.size())

print("El offset es", tensor1_transpose.storage_offset())

print("El stride es", tensor1_transpose.stride())

tensor([[2., 4., 2.],
        [1., 3., 1.]])
 2.0
 1.0
 4.0
 3.0
 2.0
 1.0
¿Es el tensor contiguo? False 
El tamaño (size) es torch.Size([2, 3]) 
El offset es 0 El stride es (1, 2) 

Como podemos ver, el almacenamiento interno del tensor transpuesto con dos filas y tres columnas no cambia pero sí que cambian los metadatos de vista. Ahora el tensor no es contiguo y el stride es (1,2), por lo que necesitamos avanzar (1,2) posiciones en el almacenamiento para avanzar una posición en la (fila/columna).