KeondoPark

Code refactoring WIP - remove redundant code blocks

2408d2e 5 months ago

3.47 kB

# 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

	# 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_sizegrid_size, embed_dim] or [1+grid_sizegrid_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