| """ |
| 2D sine-cosine positional embeddings. |
| |
| STP uses a *fixed* sincos table (the checkpoint stores it as a non-trainable |
| buffer), so these tables must be reproduced exactly at initialisation time. |
| """ |
| import numpy as np |
|
|
|
|
| def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False): |
| """2D sincos table for a `grid_size` x `grid_size` patch grid. |
| |
| Returns [grid_size**2, embed_dim], or [1 + grid_size**2, embed_dim] when |
| `cls_token=True` (a zero row is prepended for the [CLS] token). |
| """ |
| grid_h = np.arange(grid_size, dtype=np.float32) |
| grid_w = np.arange(grid_size, dtype=np.float32) |
| grid = np.meshgrid(grid_w, grid_h) |
| grid = np.stack(grid, axis=0) |
| grid = grid.reshape([2, 1, grid_size, grid_size]) |
|
|
| pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid) |
| if cls_token: |
| pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) |
| return pos_embed |
|
|
|
|
| def get_2d_sincos_pos_embed_from_grid(embed_dim, grid): |
| assert embed_dim % 2 == 0, 'embed_dim must be even to split between h and w' |
|
|
| |
| emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) |
| emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) |
| return np.concatenate([emb_h, emb_w], axis=1) |
|
|
|
|
| def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): |
| """1D sincos table: `pos` of shape (M,) -> (M, embed_dim).""" |
| assert embed_dim % 2 == 0 |
|
|
| |
| omega = np.arange(embed_dim // 2, dtype=np.float64) |
| omega /= embed_dim / 2. |
| omega = 1. / 10000 ** omega |
|
|
| pos = pos.reshape(-1) |
| out = np.einsum('m,d->md', pos, omega) |
|
|
| emb_sin = np.sin(out) |
| emb_cos = np.cos(out) |
| return np.concatenate([emb_sin, emb_cos], axis=1) |
|
|
