The core concept in TensorLy-Torch is that of factorized tensors.
We provide a
FactorizedTensor class that can be used just like any PyTorch.Tensor but
provides all tensor factorization through one, simple API.
Creating factorized tensors
You can create a new factorized tensor easily:
The signature is:
factorized_tensor = FactorizedTensor.new(shape, rank, factorization)
For instance, to create a tensor in Tucker form, that has half the parameters of a dense (non-factorized) tensor of the same shape, you would simply write:
tucker_tensor = FactorizedTensor.new(shape, rank=0.5, factorization='tucker')
Since TensorLy-Torch builds on top of TensorLy, it also comes with tensor decomposition out-of-the-box. To initialize a factorized tensor in CP (Canonical-Polyadic) form, also known as Parafac, or Kruskal tensor, with 1/10th of the parameters, you can simply write:
cp_tensor = FactorizedTensor.new(dense_tensor, rank=0.1, factorization='CP')
Manipulating factorized tensors
The first thing you want to do, if you created a new tensor from scratch (by using the
new method), is to initialize it,
e.g. so that the element of the reconstruction approximately follow a Gaussian distribution:
You can even use PyTorch’s functions! This works:
from torch.nn import init init.kaiming_normal(cp_tensor)
Finally, you can index tensors directly in factorized form, which will return another factorized tensor, whenever possible!
>>> cp_tensor[:2, :2] CPTensor(shape=(2, 2, 2), rank=2)
If not possible, a dense tensor will be returned:
>>> cp_tensor[2, 3, 1] tensor(0.0250, grad_fn=<SumBackward0>)
Note how, above, indexing tracks gradients as well!
In addition to tensor in factorized forms, TensorLy-Torch provides out-of-the-box for Tensorized tensors. The most common case is that of tensorized matrices, where a matrix is first tensorized, i.e. reshaped into a higher-order tensor which is then decomposed and stored in factorized form.
A commonly used tensorized tensor is the tensor-train matrix (also known as Matrix-Product Operator in quantum physics), or, in general, Block-TT.
You can create one in TensorLy-Torch, from a matrix, just as easily as a regular tensor, using the
with the following signature:
TensorizedTensor.from_matrix(matrix, tensorized_row_shape, tensorized_column_shape, rank)
where tensorized_row_shape and tensorized_column_shape indicate the shape to which to tensorize the row and column size of the given matrix. For instance, if you have a matrix of size 16x21, you could use tensorized_row_shape=(4, 4) and tensorized_column_shape=(3, 7).
In general, you can tensorize any tensor, not just matrices, even with batched modes (dimensions)!
tensorized_tensor = TensorizedTensor.new(tensorized_shape, rank, factorization)
tensorized_shape is a nested tuple, in which an int represents a batched mode, and a tuple a tensorized mode.
For instance, a batch of 5 matrices of size 16x21 could be tensorized into a batch of 5 tensorized matrices of size (4x4)x(3x7), in the BlockTT form. In code, you would do this using
tensorized_tensor = TensorizedTensor.from_tensor(tensor, (5, (4, 4), (3, 7)), rank=0.7, factorization='BlockTT')
You can of course tensorize any size tensors, e.g. a batch of 5 matrices of size 8x27 can be tensorized into:
>>> ftt = tltorch.TensorizedTensor.new((5, (2, 2, 2), (3, 3, 3)), rank=0.5, factorization='BlockTT')
This returns a tensorized tensor, stored in decomposed form: >>> ftt BlockTT(shape=[5, 8, 27], tensorized_shape=(5, (2, 2, 2), (3, 3, 3)), rank=[1, 20, 20, 1])
As for factorized tensors, you can directly index them:
>>> ftt BlockTT(shape=[8, 27], tensorized_shape=[(2, 2, 2), (3, 3, 3)], rank=[1, 20, 20, 1])
>>> ftt[0, :2, :2] tensor([[-0.0009, 0.0004], [ 0.0007, 0.0003]], grad_fn=<SqueezeBackward0>)
Again, notice that gradients are tracked and all operations on factorized and tensorized tensors are back-propagatable!