cancel
Showing results for 
Show  only  | Search instead for 
Did you mean: 

Vitis AI - How tensors allow for efficent memory usage

gguasti
Xilinx Employee
Xilinx Employee
5 0 1,015

In data manipulation, it is typical to reshape or reorder the original data and create multiple copies. At any new step, a new copy is created. As the program grows, so does the occupied memory, and I seldom think to worry about this issue until an Out Of Memory error happens.

The amazing thing about tensors is that multiple tensors can refer to the same storage (a contiguous chunk of memory containing a given type of numbers). This is managed by torch.storage.

Each tensor has the .storage property which shows us the tensor content as it is stored in the memory.

In the next article I will write about the even more amazing Tensor property of tracking the ancestors operations, but here I will mostly focus on memory optimization.

Note: the new Vitis AI 1.2 release will be the first to support PyTorch. This article celebrates the addition of support for the popular framework with an example in the Jupyter Notebook format dedicated entirely to PyTorch.

 

In [1]:
Out[1]:
tensor([[4, 1, 6],
        [0, 8, 8],
        [1, 2, 1],
        [0, 5, 7],
        [0, 0, 7]])
 
In [2]:
Out[2]:
 4
 1
 6
 0
 8
 8
 1
 2
 1
 0
 5
 7
 0
 0
 7
[torch.LongStorage of size 15]
 
In [3]:
 
Out[3]:
torch.Size([5, 3])
 

We might need to transpose and flatten the original "a" tensor.

Why should we waste double the memory for the same data, even with a different shape?

In [4]:
Out[4]:
tensor([[4, 0, 1, 0, 0],
        [1, 8, 2, 5, 0],
        [6, 8, 1, 7, 7]])
 

a and b are indeed tensors pointing to the same storage.

They are represented differently because we tell them to read from the storage in a different order, using the stride function.

b has a stride equal to (1,3) meaning that when reading the storage it must jump to the next row every 1 element, and jump to next column every 3 elements.

In [5]:
Out[5]:
((1, 3), (3, 1))
 

We can access the data from a, b, or directly from the original storage.

However if we access from the storage the value read is no longer a tensor.

In [6]:
Out[6]:
(tensor(8), tensor(8), 8, 8)
 

Now, one odd thing that drove me crazy was when I found that my tensors magically changed in value by themselves:

When changing a, b is also modified.

In [7]:
Out[7]:
tensor(10)
 

This happens because from a memory point of view tensors are an ordered representation of a storage.

Two tensors generated from the same storage are not independent and I must remember that every time I change one tensor, all other tensors pointing to the same storage are also modified.

An efficient usage of the memory has its drawbacks!

Subsets
 

Subset of the original data still make an efficient usage of the memory.

The new tensor will still point to a subset of the original storage.

In [8]:
Out[8]:
tensor([[10,  1],
        [ 0,  8]])
 
In [9]:
Out[9]:
tensor([[77,  1,  6],
        [ 0,  8,  8],
        [ 1,  2,  1],
        [ 0,  5,  7],
        [ 0,  0,  7]])
 
Inplace operations
 

inplace operations are functions which are able to operate directly on the storage without creating a copy of the tensor. They usually have a recognizable name with a trailing underscore.

 
In [10]:
Out[10]:
tensor([[0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0]])
 

Tensor cloning

 

When we really need an independent new tensor, we can clone it.

A new storage will also be created.

 
In [11]:
 
Out[11]:
tensor([[55,  0,  0],
        [ 0,  0,  0],
        [ 0,  0,  0],
        [ 0,  0,  0],
        [ 0,  0,  0]])
 
In [12]:
Out[12]:
tensor([[0, 0, 0],
        [0, 0, 0],
        [0, 0, 0],
        [0, 0, 0],
        [0, 0, 0]])
 

Reshuffle the storage for a contiguous tensor

Some functions only work on contiguous tensors.

When we transposed a, we generated a new tensor by allocating matrix values in b from the storage that are not contiguous.

In [13]:
 
Out[13]:
True
 
In [14]:
Out[14]:
False
 

We can make b contiguous, but this will generate a new reshuffled storage from b, and a and b will be forever independent:

In [15]:
Out[15]:
tensor(0)
 
In [16]:
Out[16]:
True
 
In [17]:
Out[17]:
True