Spaces:
Sleeping
Sleeping
File size: 1,311 Bytes
6e9c433 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
import numpy as np
from torch import Tensor
import matplotlib.pyplot as plt
from s_multimae.model.multimae import build_2d_sincos_posemb
def visualize_2d_posemb():
NH, NW = 14, 14
dim_tokens = 768
colors = [
"Greys",
"Purples",
"Blues",
"Greens",
"Oranges",
"Reds",
"YlOrBr",
"YlOrRd",
"OrRd",
"PuRd",
"RdPu",
"BuPu",
"GnBu",
"PuBu",
"YlGnBu",
"PuBuGn",
"BuGn",
"YlGn",
]
pos_emb: Tensor = build_2d_sincos_posemb(NH, NW, dim_tokens)
pos_emb_numpy: np.ndarray = (
pos_emb.squeeze(0).permute(1, 2, 0).numpy()
) # 14 x 14 x 768
x = np.linspace(0, NH - 1, NH)
y = np.linspace(0, NW - 1, NW)
X, Y = np.meshgrid(x, y)
for color, i in zip(colors, range(0, pos_emb_numpy.shape[2], 100)):
ax = plt.axes(projection="3d")
Z = pos_emb_numpy[:, :, i]
# plt.imshow(Z, cmap='viridis')
# plt.savefig(f'posemb_visualization/test_{i}.png')
ax.plot_surface(
X,
Y,
Z,
# rstride=1, cstride=1,
cmap="viridis",
edgecolor="none",
)
plt.show()
plt.savefig(f"posemb_visualization/test_{i}.png")
|