[Python] efficient matrix inversion/multiplication with multiple batch dimensions in pytorch

In my framework, I am having an outer loop (here mocked by the variable n) and inside the loop body I have to perform matrix inversions/multiplications for multiple batch dimensions. I observed that manually looping over the batch dimensions and computing the inverse is measurable faster than passing all batches to the torch.linalg.inv function. Similar statements can be made for computing matrix multiplications using torch.einsum function.

My expectation was that passing all batches at once performs better than for-loops. Any ideas/explanations/recommendations here?

Profiling output of function inverse_batch:

ncalls tottime percall cumtime percall filename:lineno(function)
100 2.112 0.021 2.112 0.021 {built-in method torch._C._linalg.linalg_inv}
1 0.000 0.000 2.112 2.112 mwe.py:5(inverse_batch)
1 0.000 0.000 2.112 2.112 {built-in method builtins.exec}
1 0.000 0.000 2.112 2.112 <string>:1(<module>)
1 0.000 0.000 0.000 0.000 {method 'disable' of '_lsprof.Profiler' objects}


Profiling output of function inverse_loop:

ncalls tottime percall cumtime percall filename:lineno(function)
8000 0.207 0.000 0.207 0.000 {built-in method torch._C._linalg.linalg_inv}
1 0.022 0.022 0.229 0.229 mwe.py:9(inverse_loop)
1 0.000 0.000 0.000 0.000 {method 'view' of 'torch._C.TensorBase' objects}
1 0.000 0.000 0.229 0.229 {built-in method builtins.exec}
1 0.000 0.000 0.229 0.229 <string>:1(<module>)
1 0.000 0.000 0.000 0.000 {method 'disable' of '_lsprof.Profiler' objects}


Code:

import torch
import cProfile
import pstats

def inverse_batch(tensors, n):
for i in range(n):
torch.linalg.inv(tensors)

def inverse_loop(tensors, n):
tensors = tensors.view(-1, 3, 3)
for i in range(n):
for j in range(10 * 8):
torch.linalg.inv(tensors[j])


# Create a batch of tensors
tensors = torch.randn(10, 8, 3, 3, dtype = torch.double) # Shape: (10, 8, 3, 3)

# Profile code
n = 100 # Dummy outer loop variable
cProfile.run('inverse_batch(tensors, n)', 'profile_output')
stats = pstats.Stats('profile_output')
stats.strip_dirs().sort_stats('tottime').print_stats()

Continue reading...