File size: 3,468 Bytes
35f60b9 2408d2e 35f60b9 | 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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# Position embedding utils
# --------------------------------------------------------
import numpy as np
import torch
# --------------------------------------------------------
# 2D sine-cosine position embedding
# References:
# Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py
# MoCo v3: https://github.com/facebookresearch/moco-v3
# --------------------------------------------------------
def get_1d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
"""
grid_size: int of the grid height and width
return:
pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
"""
# grid = grid.reshape([2, 1, grid_size, grid_size])
grid = np.arange(grid_size, dtype=np.float32)
grid = grid.reshape([1, 1, grid_size])
pos_embed = get_1d_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_1d_sincos_pos_embed_from_grid(embed_dim, pos):
"""
embed_dim: output dimension for each position
pos: a list of positions to be encoded: size (M,)
out: (M, D)
"""
assert embed_dim % 2 == 0
omega = np.arange(embed_dim // 2, dtype=np.float32)
omega /= embed_dim / 2.
omega = 1. / 10000**omega # (D/2,)
pos = pos.reshape(-1) # (M,)
out = np.einsum('m,d->md', pos, omega) # (M, D/2), outer product
emb_sin = np.sin(out) # (M, D/2)
emb_cos = np.cos(out) # (M, D/2)
emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)
return emb
def flip_padded_sequences(batch, mask):
"""
Flips the valid part of each sequence in a padded batch, keeping pads at the end.
Args:
batch (Tensor): Input tensor of size (B, L).
lengths (Tensor, optional): Lengths of size (B,). If None, calc from mask.
padding_value (int): Value used for padding.
Returns:
Tensor: The flipped batch.
"""
B, L, D = batch.shape
device = batch.device
# 1. Get Sequence Lengths
lengths = mask.sum(dim=1).long()
# 2. Generate Index Grid
# Shape: (L,) -> [0, 1, 2, ..., L-1]
range_tensor = torch.arange(L, device=device)
# Expand to (B, L)
# [[0, 1, ...], [0, 1, ...], ...]
idx_grid = range_tensor.unsqueeze(0).expand(B, L)
# 3. Calculate "Inverse" Indices for valid positions
# If len=3, indices [0, 1, 2] should become [2, 1, 0]
# Formula: flipped_index = (Length - 1) - current_index
lens_expanded = lengths.unsqueeze(1)
flipped_indices = (lens_expanded - 1) - idx_grid
# 4. Create the final Gather Indices
# We only want to use the flipped indices where the data is valid.
# Where it is padding, we keep the original index (so we just gather the 0s back).
mask = idx_grid < lens_expanded
# Use torch.where: if valid, use flipped_indices, else use original idx_grid
gather_indices = torch.where(mask, flipped_indices, idx_grid)
# 5. Gather
result = torch.gather(batch, 1, gather_indices.unsqueeze(-1).repeat(1, 1, D))
return result
|