Tensor Train layers for neural networks

In this section, the TT layers are introduced.

Imports:

import torch as tn
import torch.nn as nn
import datetime
try:
    import torchtt as tntt
except:
    print('Installing torchTT...')
    %pip install git+https://github.com/ion-g-ion/torchTT
    import torchtt as tntt

We consider a linear layer \(\mathcal{LTT}(\mathsf{x}) = \mathsf{Wx}+\mathsf{b}\) acting on a tensor input \(\mathsf{x}\) of shape \(n_1 \times \cdots \times n_d\) and returning a tensor of shape \(m_1\times\cdots\times m_d\). The corresponding weight matrix \(\mathsf{W}\) would have the shape \((m_1\times\cdots\times m_d) \times (n_1 \times \cdots \times n_d)\). The goal is to represent the weights tensor operator in TT format and perform the learning with respect tot the cores of the TT decomposition (ranks have to be fixed a priori). Due to the AD functionality of torchtt, the gradient with respect tot the cores can be computed for any network structure. TT layers can be added using torchtt.nn.LinearLayerTT() class.

In the following, a neural netywork with 3 hidden layers and one linear layer is created. The shapes of the individual layers are

\(\mathbb{R}^{16} \times\mathbb{R}^{16} \times\mathbb{R}^{16} \times\mathbb{R}^{16} \underset{}{\longrightarrow} \mathbb{R}^8 \times\mathbb{R}^8 \times\mathbb{R}^8 \times\mathbb{R}^8 \underset{}{\longrightarrow} \mathbb{R}^4 \times\mathbb{R}^4 \times\mathbb{R}^4 \times\mathbb{R}^4 \underset{}{\longrightarrow} \mathbb{R}^2 \times\mathbb{R}^4 \times\mathbb{R}^2 \times\mathbb{R}^4 \underset{}{\longrightarrow} \mathbb{R}^{10}\).

class BasicTT(nn.Module):
    def __init__(self):
        super().__init__()
        self.ttl1 = tntt.nn.LinearLayerTT([16,16,16,16], [8,8,8,8], [1,3,3,3,1])
        self.ttl2 = tntt.nn.LinearLayerTT([8,8,8,8], [4,4,4,4], [1,2,2,2,1])
        self.ttl3 = tntt.nn.LinearLayerTT([4,4,4,4], [2,4,2,4], [1,2,2,2,1])
        self.linear = nn.Linear(64, 10, dtype = tn.float32)

    def forward(self, x):
        x = self.ttl1(x)
        x = tn.relu(x)
        x = self.ttl2(x)
        x = tn.relu(x)
        x = self.ttl3(x)
        x = tn.relu(x)
        x = tn.reshape(x,[-1,64])
        return self.linear(x)

Create the model and print the number of trainable parameters as well as the model structure.

model = BasicTT()
print('Number of trainable parameters:', len(list(model.parameters())))
print(model)

Number of trainable parameters: 17
BasicTT(
  (ttl1): LinearLayerTT(
    (cores): ParameterList(
        (0): Parameter containing: [torch.float32 of size 1x8x16x3]
        (1): Parameter containing: [torch.float32 of size 3x8x16x3]
        (2): Parameter containing: [torch.float32 of size 3x8x16x3]
        (3): Parameter containing: [torch.float32 of size 3x8x16x1]
    )
  )
  (ttl2): LinearLayerTT(
    (cores): ParameterList(
        (0): Parameter containing: [torch.float32 of size 1x4x8x2]
        (1): Parameter containing: [torch.float32 of size 2x4x8x2]
        (2): Parameter containing: [torch.float32 of size 2x4x8x2]
        (3): Parameter containing: [torch.float32 of size 2x4x8x1]
    )
  )
  (ttl3): LinearLayerTT(
    (cores): ParameterList(
        (0): Parameter containing: [torch.float32 of size 1x2x4x2]
        (1): Parameter containing: [torch.float32 of size 2x4x4x2]
        (2): Parameter containing: [torch.float32 of size 2x2x4x2]
        (3): Parameter containing: [torch.float32 of size 2x4x4x1]
    )
  )
  (linear): Linear(in_features=64, out_features=10, bias=True)
)

A random input is created and passed as argument to the model. Batch evaluation is also possible by extending the dimensionality of the input before the leading mode.

input = tn.rand((16,16,16,16), dtype = tn.float32)
pred = model.forward(input)

input_batch = tn.rand((1000,16,16,16,16), dtype = tn.float32)
label_batch = tn.rand((1000,10), dtype = tn.float32)
%time pred = model.forward(input_batch)

CPU times: user 3.61 s, sys: 2.22 s, total: 5.84 s
Wall time: 450 ms

The obtained network can be trained similarily to other torch models. A loss function together with an optimizer are defined.

criterion = nn.CrossEntropyLoss()
optimizer = tn.optim.Adam(model.parameters(), lr = 0.001)

A training loop is executed to exemplify the training parameters update procedure. An example where a true dataset is used is presented here.

for epoch in range(5):

    optimizer.zero_grad()

    outputs = model(input_batch)
    loss = criterion(outputs, label_batch)
    loss.backward()
    optimizer.step()

    # print statistics
    print('Epoch %d, loss %e'%(epoch+1,loss.item()))


print('Finished Training')
%time plm = model(input_batch)

Epoch 1, loss 1.196611e+01
Epoch 2, loss 1.182479e+01
Epoch 3, loss 1.174133e+01
Epoch 4, loss 1.169423e+01
Epoch 5, loss 1.166731e+01
Finished Training
CPU times: user 3.85 s, sys: 1.79 s, total: 5.64 s
Wall time: 499 ms

If the GPU is available, the model can be run on it to get a speedup (should be run 2 times to see the speedup due to CUDA warm-up).

if tn.cuda.is_available():
    model_gpu = BasicTT().cuda()
    input_batch_gpu = tn.rand((400,16,16,16,16)).cuda()

    input_batch = tn.rand((400,16,16,16,16))
    tme = datetime.datetime.now()
    pred = model.forward(input_batch)
    tme = datetime.datetime.now() - tme
    print('Time on CPU ',tme)

    tme = datetime.datetime.now()
    pred_gpu = model_gpu.forward(input_batch_gpu).cpu()
    tme = datetime.datetime.now() - tme
    print('Time on GPU ',tme)

Time on CPU  0:00:00.232971
Time on GPU  0:00:00.003855