python - How can I make this PyTorch heatmap function faster and more efficient?

Question

I have this function that creates a sort if heatmap for 2d tensors, but it's painfully slow when using larger tensor inputs. How can I speed it up and make it more efficient?

import torch
import numpy as np
import matplotlib.pyplot as plt


def heatmap(
    tensor: torch.Tensor,
) -> torch.Tensor:
    assert tensor.dim() == 2

    def color_tensor(x: torch.Tensor) -> torch.Tensor:
        if x < 0:
            x = -x
            if x < 0.5:
                x = x * 2
                return (1 - x) * torch.tensor(
                    [0.9686, 0.9686, 0.9686]
                ) + x * torch.tensor([0.5725, 0.7725, 0.8706])
            else:
                x = (x - 0.5) * 2
                return (1 - x) * torch.tensor(
                    [0.5725, 0.7725, 0.8706]
                ) + x * torch.tensor([0.0196, 0.4431, 0.6902])
        else:
            if x < 0.5:
                x = x * 2
                return (1 - x) * torch.tensor(
                    [0.9686, 0.9686, 0.9686]
                ) + x * torch.tensor([0.9569, 0.6471, 0.5098])
            else:
                x = (x - 0.5) * 2
                return (1 - x) * torch.tensor(
                    [0.9569, 0.6471, 0.5098]
                ) + x * torch.tensor([0.7922, 0.0000, 0.1255])

    return torch.stack(
        [torch.stack([color_tensor(x) for x in t]) for t in tensor]
    ).permute(2, 0, 1)

x = torch.randn(3,3)
x = x / x.max()
x_out = heatmap(x)

x_out = (x_out.permute(1, 2, 0) * 255).numpy()
plt.imshow(x_out.astype(np.uint8))
plt.axis("off")
plt.show()

An example of the output:

Answer 1

You need to get rid of ifs and the for loop and make a vectorized function. To do that, you can use masks and calculate all in one. Here it is:

def heatmap(tensor: torch.Tensor) -> torch.Tensor:
    assert tensor.dim() == 2

    # We're expanding to create one more dimension, for mult. to work.
    xt = x.expand((3, x.shape[0], x.shape[1])).permute(1, 2, 0)

    # this part is the mask: (xt >= 0) * (xt < 0.5) ...
    # ... the rest is the original function translated
    color_tensor = (
        (xt >= 0) * (xt < 0.5) * ((1 - xt * 2) * torch.tensor([0.9686, 0.9686, 0.9686]) + xt * 2 * torch.tensor([0.9569, 0.6471, 0.5098]))
        +
        (xt >= 0) * (xt >= 0.5) * ((1 - (xt - 0.5) * 2) * torch.tensor([0.9569, 0.6471, 0.5098]) + (xt - 0.5) * 2 * torch.tensor([0.7922, 0.0000, 0.1255]))
        +
        (xt < 0) * (xt > -0.5) * ((1 - (-xt * 2)) * torch.tensor([0.9686, 0.9686, 0.9686]) + (-xt * 2) * torch.tensor([0.5725, 0.7725, 0.8706]))
        +
        (xt < 0) * (xt <= -0.5) * ((1 - (-xt - 0.5) * 2) * torch.tensor([0.5725, 0.7725, 0.8706]) + (-xt - 0.5) * 2 * torch.tensor([0.0196, 0.4431, 0.6902]))
    ).permute(2, 0, 1)
    
    return color_tensor