#python #arrays #numpy #pytorch #tensor
#python #массивы #numpy #pytorch #тензор
Вопрос:
У меня есть эта функция, которая создает тепловую карту sort if для 2d-тензоров, но она работает очень медленно при использовании больших тензорных входов. Как я могу ускорить ее и сделать более эффективной?
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()
Пример вывода:
Ответ №1:
Вам нужно избавиться от if
s и цикла for и создать векторизованную функцию. Для этого вы можете использовать маски и вычислять все в одном. Вот оно:
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
Комментарии:
1. Ваша векторизованная версия работает намного лучше, чем мой код. Спасибо, что помогли мне!