|
|
|
|
|
|
|
|
|
|
|
|
|
from typing import Optional, Tuple, Union |
|
|
|
import math |
|
import torch |
|
import torch.nn.functional as F |
|
import torch.utils.checkpoint |
|
from einops import rearrange |
|
from timm.models.layers import DropPath |
|
from torch import nn |
|
from transformers.activations import ACT2FN |
|
from transformers.modeling_outputs import (BaseModelOutput, |
|
BaseModelOutputWithPooling) |
|
from transformers.modeling_utils import PreTrainedModel |
|
from transformers.utils import logging |
|
|
|
from .configuration_intern_vit import InternVisionConfig |
|
|
|
try: |
|
from flash_attn.bert_padding import pad_input, unpad_input |
|
from flash_attn.flash_attn_interface import \ |
|
flash_attn_varlen_qkvpacked_func, flash_attn_varlen_func |
|
has_flash_attn = True |
|
except: |
|
print('FlashAttention2 is not installed.') |
|
has_flash_attn = False |
|
|
|
logger = logging.get_logger(__name__) |
|
|
|
|
|
class FlashAttention(nn.Module): |
|
"""Implement the scaled dot product attention with softmax. |
|
Arguments |
|
--------- |
|
softmax_scale: The temperature to use for the softmax attention. |
|
(default: 1/sqrt(d_keys) where d_keys is computed at |
|
runtime) |
|
attention_dropout: The dropout rate to apply to the attention |
|
(default: 0.0) |
|
""" |
|
|
|
def __init__(self, softmax_scale=None, attention_dropout=0.0, device=None, dtype=None): |
|
super().__init__() |
|
self.softmax_scale = softmax_scale |
|
self.dropout_p = attention_dropout |
|
|
|
def forward(self, qkv, key_padding_mask=None, causal=False, cu_seqlens=None, |
|
max_s=None, need_weights=False): |
|
"""Implements the multihead softmax attention. |
|
Arguments |
|
--------- |
|
qkv: The tensor containing the query, key, and value. (B, S, 3, H, D) if key_padding_mask is None |
|
if unpadded: (nnz, 3, h, d) |
|
key_padding_mask: a bool tensor of shape (B, S) |
|
""" |
|
assert not need_weights |
|
assert qkv.dtype in [torch.float16, torch.bfloat16] |
|
assert qkv.is_cuda |
|
|
|
if cu_seqlens is None: |
|
batch_size = qkv.shape[0] |
|
seqlen = qkv.shape[1] |
|
if key_padding_mask is None: |
|
qkv = rearrange(qkv, 'b s ... -> (b s) ...') |
|
max_s = seqlen |
|
cu_seqlens = torch.arange(0, (batch_size + 1) * seqlen, step=seqlen, dtype=torch.int32, |
|
device=qkv.device) |
|
output = flash_attn_varlen_qkvpacked_func( |
|
qkv, cu_seqlens, max_s, self.dropout_p if self.training else 0.0, |
|
softmax_scale=self.softmax_scale, causal=causal |
|
) |
|
output = rearrange(output, '(b s) ... -> b s ...', b=batch_size) |
|
else: |
|
nheads = qkv.shape[-2] |
|
x = rearrange(qkv, 'b s three h d -> b s (three h d)') |
|
x_unpad, indices, cu_seqlens, max_s = unpad_input(x, key_padding_mask) |
|
x_unpad = rearrange(x_unpad, 'nnz (three h d) -> nnz three h d', three=3, h=nheads) |
|
output_unpad = flash_attn_varlen_qkvpacked_func( |
|
x_unpad, cu_seqlens, max_s, self.dropout_p if self.training else 0.0, |
|
softmax_scale=self.softmax_scale, causal=causal |
|
) |
|
output = rearrange(pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'), |
|
indices, batch_size, seqlen), |
|
'b s (h d) -> b s h d', h=nheads) |
|
else: |
|
assert max_s is not None |
|
output = flash_attn_varlen_qkvpacked_func( |
|
qkv, cu_seqlens, max_s, self.dropout_p if self.training else 0.0, |
|
softmax_scale=self.softmax_scale, causal=causal |
|
) |
|
|
|
return output, None |
|
|
|
|
|
class InternRMSNorm(nn.Module): |
|
def __init__(self, hidden_size, eps=1e-6): |
|
super().__init__() |
|
self.weight = nn.Parameter(torch.ones(hidden_size)) |
|
self.variance_epsilon = eps |
|
|
|
def forward(self, hidden_states): |
|
input_dtype = hidden_states.dtype |
|
hidden_states = hidden_states.to(torch.float32) |
|
variance = hidden_states.pow(2).mean(-1, keepdim=True) |
|
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) |
|
return self.weight * hidden_states.to(input_dtype) |
|
|
|
|
|
try: |
|
from apex.normalization import FusedRMSNorm |
|
|
|
InternRMSNorm = FusedRMSNorm |
|
|
|
logger.info('Discovered apex.normalization.FusedRMSNorm - will use it instead of InternRMSNorm') |
|
except ImportError: |
|
|
|
pass |
|
except Exception: |
|
logger.warning('discovered apex but it failed to load, falling back to InternRMSNorm') |
|
pass |
|
|
|
|
|
NORM2FN = { |
|
'rms_norm': InternRMSNorm, |
|
'layer_norm': nn.LayerNorm, |
|
} |
|
|
|
|
|
class InternVisionEmbeddings(nn.Module): |
|
def __init__(self, config: InternVisionConfig): |
|
super().__init__() |
|
self.config = config |
|
self.embed_dim = config.hidden_size |
|
self.image_size = config.image_size |
|
self.patch_size = config.patch_size |
|
|
|
self.class_embedding = nn.Parameter( |
|
torch.randn(1, 1, self.embed_dim), |
|
) |
|
|
|
self.patch_embedding = nn.Conv2d( |
|
in_channels=3, out_channels=self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size |
|
) |
|
|
|
self.num_patches = (self.image_size // self.patch_size) ** 2 |
|
self.num_positions = self.num_patches + 1 |
|
|
|
self.position_embedding = nn.Parameter(torch.randn(1, self.num_positions, self.embed_dim)) |
|
|
|
def _get_pos_embed(self, pos_embed, H, W): |
|
target_dtype = pos_embed.dtype |
|
pos_embed = pos_embed.float().reshape( |
|
1, self.image_size // self.patch_size, self.image_size // self.patch_size, -1).permute(0, 3, 1, 2) |
|
pos_embed = F.interpolate(pos_embed, size=(H, W), mode='bicubic', align_corners=False). \ |
|
reshape(1, -1, H * W).permute(0, 2, 1).to(target_dtype) |
|
return pos_embed |
|
|
|
def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor: |
|
target_dtype = self.patch_embedding.weight.dtype |
|
patch_embeds = self.patch_embedding(pixel_values) |
|
batch_size, _, height, width = patch_embeds.shape |
|
patch_embeds = patch_embeds.flatten(2).transpose(1, 2) |
|
class_embeds = self.class_embedding.expand(batch_size, 1, -1).to(target_dtype) |
|
embeddings = torch.cat([class_embeds, patch_embeds], dim=1) |
|
position_embedding = torch.cat([ |
|
self.position_embedding[:, :1, :], |
|
self._get_pos_embed(self.position_embedding[:, 1:, :], height, width) |
|
], dim=1) |
|
embeddings = embeddings + position_embedding.to(target_dtype) |
|
return embeddings |
|
|
|
|
|
class InternAttention(nn.Module): |
|
"""Multi-headed attention from 'Attention Is All You Need' paper""" |
|
|
|
def __init__(self, config: InternVisionConfig): |
|
super().__init__() |
|
self.config = config |
|
self.embed_dim = config.hidden_size |
|
self.num_heads = config.num_attention_heads |
|
self.use_flash_attn = config.use_flash_attn and has_flash_attn |
|
if config.use_flash_attn and not has_flash_attn: |
|
print('Warning: Flash Attention is not available, use_flash_attn is set to False.') |
|
self.head_dim = self.embed_dim // self.num_heads |
|
if self.head_dim * self.num_heads != self.embed_dim: |
|
raise ValueError( |
|
f'embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:' |
|
f' {self.num_heads}).' |
|
) |
|
|
|
self.scale = self.head_dim ** -0.5 |
|
self.qkv = nn.Linear(self.embed_dim, 3 * self.embed_dim, bias=config.qkv_bias) |
|
self.attn_drop = nn.Dropout(config.attention_dropout) |
|
self.proj_drop = nn.Dropout(config.dropout) |
|
|
|
self.qk_normalization = config.qk_normalization |
|
|
|
if self.qk_normalization: |
|
self.q_norm = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps) |
|
self.k_norm = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps) |
|
|
|
if self.use_flash_attn: |
|
self.inner_attn = FlashAttention(attention_dropout=config.attention_dropout) |
|
self.proj = nn.Linear(self.embed_dim, self.embed_dim) |
|
|
|
def _naive_attn(self, x): |
|
B, N, C = x.shape |
|
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) |
|
q, k, v = qkv.unbind(0) |
|
|
|
if self.qk_normalization: |
|
B_, H_, N_, D_ = q.shape |
|
q = self.q_norm(q.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2) |
|
k = self.k_norm(k.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2) |
|
|
|
attn = ((q * self.scale) @ k.transpose(-2, -1)) |
|
attn = attn.softmax(dim=-1) |
|
attn = self.attn_drop(attn) |
|
|
|
x = (attn @ v).transpose(1, 2).reshape(B, N, C) |
|
x = self.proj(x) |
|
x = self.proj_drop(x) |
|
return x |
|
|
|
def _flash_attn(self, x, key_padding_mask=None, need_weights=False): |
|
qkv = self.qkv(x) |
|
qkv = rearrange(qkv, 'b s (three h d) -> b s three h d', three=3, h=self.num_heads) |
|
|
|
if self.qk_normalization: |
|
q, k, v = qkv.unbind(2) |
|
q = self.q_norm(q.flatten(-2, -1)).view(q.shape) |
|
k = self.k_norm(k.flatten(-2, -1)).view(k.shape) |
|
qkv = torch.stack([q, k, v], dim=2) |
|
|
|
context, _ = self.inner_attn( |
|
qkv, key_padding_mask=key_padding_mask, need_weights=need_weights, causal=False |
|
) |
|
outs = self.proj(rearrange(context, 'b s h d -> b s (h d)')) |
|
outs = self.proj_drop(outs) |
|
return outs |
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
|
x = self._naive_attn(hidden_states) if not self.use_flash_attn else self._flash_attn(hidden_states) |
|
return x |
|
|
|
|
|
class InternMLP(nn.Module): |
|
def __init__(self, config: InternVisionConfig): |
|
super().__init__() |
|
self.config = config |
|
self.act = ACT2FN[config.hidden_act] |
|
self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) |
|
self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) |
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
|
hidden_states = self.fc1(hidden_states) |
|
hidden_states = self.act(hidden_states) |
|
hidden_states = self.fc2(hidden_states) |
|
return hidden_states |
|
|
|
|
|
def generate_batch_temporal_mask(split_sizes, device='cpu'): |
|
""" |
|
generate the temporal (padding) mask of a batch |
|
Args: |
|
split_sizes: List[num frames] |
|
Returns: |
|
temporal_mask: BoolTensor(B, T), `True` means taking, `False` means padding |
|
""" |
|
B, T = len(split_sizes), max(split_sizes) |
|
split_sizes = torch.tensor(split_sizes, dtype=torch.long, device=device) |
|
temporal_idx = torch.arange(T, dtype=torch.long, device=device)[None].repeat((B, 1)) |
|
temporal_mask = temporal_idx < split_sizes[:, None] |
|
return temporal_mask |
|
|
|
def concat_batch_frames(images, split_sizes=None, temporal_mask=None): |
|
""" |
|
B, T, L, D -> concat(T), L, D |
|
""" |
|
if temporal_mask is None: |
|
assert split_sizes is not None |
|
temporal_mask = generate_batch_temporal_mask(split_sizes, device=images.device) |
|
return images[temporal_mask] |
|
|
|
def stack_batch_frames(images, split_sizes, return_mask=False): |
|
""" |
|
concat(T), L, D -> B, T, L, D |
|
""" |
|
B, T = len(split_sizes), max(split_sizes) |
|
images_stack = images.new_zeros((B, T, *images.shape[1:])) |
|
temporal_mask = generate_batch_temporal_mask(split_sizes, device=images.device) |
|
images_stack[temporal_mask] = images |
|
if return_mask: |
|
return images_stack, temporal_mask |
|
return images_stack |
|
|
|
def temporal_idx_abs_to_rel(temporal_idx, split_sizes): |
|
stacked_temporal_idx = stack_batch_frames(temporal_idx, split_sizes) |
|
length = stacked_temporal_idx.max(dim=-1, keepdim=True)[0] |
|
length = length.clip(min=1) |
|
rel_temporal_idx = stacked_temporal_idx.float() / length.float() |
|
rel_temporal_idx = concat_batch_frames(rel_temporal_idx, split_sizes) |
|
return rel_temporal_idx |
|
|
|
|
|
def get_timestep_embedding( |
|
timesteps: torch.Tensor, |
|
embedding_dim: int, |
|
flip_sin_to_cos: bool = False, |
|
downscale_freq_shift: float = 1, |
|
scale: float = 1, |
|
max_period: int = 10000, |
|
): |
|
""" |
|
This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. |
|
|
|
Args |
|
timesteps (torch.Tensor): |
|
a 1-D Tensor of N indices, one per batch element. These may be fractional. |
|
embedding_dim (int): |
|
the dimension of the output. |
|
flip_sin_to_cos (bool): |
|
Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False) |
|
downscale_freq_shift (float): |
|
Controls the delta between frequencies between dimensions |
|
scale (float): |
|
Scaling factor applied to the embeddings. |
|
max_period (int): |
|
Controls the maximum frequency of the embeddings |
|
Returns |
|
torch.Tensor: an [N x dim] Tensor of positional embeddings. |
|
""" |
|
assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" |
|
original_dtype = timesteps.dtype |
|
|
|
half_dim = embedding_dim // 2 |
|
exponent = -math.log(max_period) * torch.arange( |
|
start=0, end=half_dim, dtype=torch.float32, device=timesteps.device |
|
) |
|
exponent = exponent / (half_dim - downscale_freq_shift) |
|
|
|
emb = torch.exp(exponent) |
|
emb = timesteps[:, None].float() * emb[None, :] |
|
|
|
|
|
emb = scale * emb |
|
|
|
|
|
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) |
|
|
|
|
|
if flip_sin_to_cos: |
|
emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) |
|
|
|
|
|
if embedding_dim % 2 == 1: |
|
emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) |
|
return emb.to(original_dtype) |
|
|
|
class Timesteps(nn.Module): |
|
def __init__(self, num_channels: int, flip_sin_to_cos: bool = False, downscale_freq_shift: float = 0, scale: int = 1): |
|
super().__init__() |
|
self.num_channels = num_channels |
|
self.flip_sin_to_cos = flip_sin_to_cos |
|
self.downscale_freq_shift = downscale_freq_shift |
|
self.scale = scale |
|
|
|
def forward(self, timesteps): |
|
t_emb = get_timestep_embedding( |
|
timesteps, |
|
self.num_channels, |
|
flip_sin_to_cos=self.flip_sin_to_cos, |
|
downscale_freq_shift=self.downscale_freq_shift, |
|
scale=self.scale, |
|
) |
|
return t_emb |
|
|
|
class AdaLayerNorm(nn.Module): |
|
def __init__( |
|
self, |
|
embedding_dim: int, |
|
conditioning_embedding_dim: int, |
|
elementwise_affine=False, |
|
eps=1e-5, |
|
bias=True, |
|
norm_type="layer_norm", |
|
zero_init=False, |
|
): |
|
super().__init__() |
|
self.silu = nn.SiLU() |
|
self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias) |
|
if zero_init: |
|
nn.init.zeros_(self.linear.weight) |
|
nn.init.zeros_(self.linear.bias) |
|
print('AdaLN zero init') |
|
if norm_type == "layer_norm": |
|
self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias) |
|
else: |
|
raise ValueError(f"unknown norm_type {norm_type}") |
|
|
|
def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor: |
|
emb = self.linear(self.silu(conditioning_embedding).to(x.dtype)) |
|
scale, shift = torch.chunk(emb, 2, dim=-1) |
|
x = self.norm(x) * (1 + scale) + shift |
|
return x |
|
|
|
|
|
class TokenTemporalAttention(nn.Module): |
|
def __init__(self, config: InternVisionConfig): |
|
super().__init__() |
|
self.config = config |
|
|
|
d_model = config.hidden_size |
|
temporal_num_heads = config.num_attention_heads |
|
self.temporal_attn = nn.MultiheadAttention(d_model, temporal_num_heads, batch_first=True) |
|
|
|
self.timestep_scale = self.config.relative_timestep_scale |
|
self.time_embed = nn.Sequential( |
|
Timesteps(num_channels=256), |
|
nn.Linear(256, d_model), |
|
nn.SiLU(), |
|
nn.Linear(d_model, d_model), |
|
) |
|
self.adaln = AdaLayerNorm(d_model, d_model, eps=config.layer_norm_eps, |
|
zero_init=self.config.temporal_adaln_zero_init) |
|
if self.config.temporal_adaln_hidden_condition: |
|
self.hidden_condition_proj = nn.Sequential( |
|
nn.Linear(d_model, d_model), |
|
nn.SiLU(), |
|
nn.Linear(d_model, d_model) |
|
) |
|
|
|
if self.config.temporal_alpha_channelwise: |
|
self.alpha_xattn = nn.Parameter(self.config.temporal_alpha_init * torch.ones(d_model), |
|
requires_grad=True) |
|
else: |
|
self.alpha_xattn = nn.Parameter(torch.tensor(self.config.temporal_alpha_init), requires_grad=True) |
|
|
|
def forward(self, |
|
hidden_states: torch.Tensor, |
|
split_sizes: Optional[list] = None, |
|
place: Optional[str] = None, |
|
temporal_id: Optional[torch.LongTensor] = None, |
|
): |
|
|
|
if self.config.use_flash_attn: |
|
return self._forward_flash_attention_2(hidden_states, split_sizes, place, temporal_id) |
|
|
|
|
|
hidden_states = stack_batch_frames(hidden_states, split_sizes) |
|
residual = hidden_states |
|
B, T, L, D = hidden_states.shape |
|
x = hidden_states.transpose(1, 2).flatten(0, 1) |
|
|
|
|
|
temporal_mask = generate_batch_temporal_mask(split_sizes, device=hidden_states.device) |
|
temporal_mask = temporal_mask.unsqueeze(1).expand(B, L, T).flatten(0, 1) |
|
if self.config.temporal_causal: |
|
attn_mask = torch.ones(T, T, dtype=torch.bool, device=hidden_states.device).tril(diagonal=0) |
|
else: |
|
attn_mask = None |
|
|
|
|
|
timestep = temporal_idx_abs_to_rel(temporal_id, split_sizes) |
|
timestep = timestep * self.timestep_scale |
|
time_condition = self.time_embed(timestep.to(hidden_states.dtype)) |
|
time_condition = stack_batch_frames(time_condition, split_sizes) |
|
time_condition = time_condition.unsqueeze(1).repeat(1, L, 1, 1).flatten(0, 1) |
|
condition = time_condition |
|
if self.config.temporal_adaln_hidden_condition: |
|
condition = condition + self.hidden_condition_proj(x) |
|
x = self.adaln(x, condition) |
|
|
|
|
|
q = k = v = x |
|
attn_mask = ~attn_mask if attn_mask is not None else None |
|
temporal_mask = ~temporal_mask |
|
|
|
attn_out = self.temporal_attn(q, k, v, attn_mask=attn_mask, key_padding_mask=temporal_mask) |
|
x = attn_out[0] |
|
|
|
|
|
x = x.view(B, L, T, D).transpose(1, 2) |
|
hidden_states = residual + x * self.alpha_xattn |
|
|
|
|
|
hidden_states = concat_batch_frames(hidden_states, split_sizes) |
|
|
|
return hidden_states |
|
|
|
def _forward_flash_attention_2(self, |
|
hidden_states: torch.Tensor, |
|
split_sizes: Optional[list] = None, |
|
place: Optional[str] = None, |
|
temporal_id: Optional[torch.LongTensor] = None, |
|
): |
|
B, T = len(split_sizes), max(split_sizes) |
|
N, L, D = hidden_states.shape |
|
residual = hidden_states |
|
hidden_states = hidden_states.transpose(0, 1).flatten(0, 1) |
|
|
|
|
|
timestep = temporal_idx_abs_to_rel(temporal_id, split_sizes) |
|
timestep = timestep * self.timestep_scale |
|
time_condition = self.time_embed(timestep.to(hidden_states.dtype)) |
|
time_condition = time_condition.unsqueeze(0).repeat(L, 1, 1).flatten(0, 1) |
|
condition = time_condition |
|
if self.config.temporal_adaln_hidden_condition: |
|
condition = condition + self.hidden_condition_proj(hidden_states) |
|
hidden_states = self.adaln(hidden_states, condition) |
|
|
|
q = k = v = hidden_states |
|
w_q, w_k, w_v = self.temporal_attn.in_proj_weight.chunk(3) |
|
b_q, b_k, b_v = self.temporal_attn.in_proj_bias.chunk(3) |
|
q = F.linear(q, w_q, b_q) |
|
k = F.linear(k, w_k, b_k) |
|
v = F.linear(v, w_v, b_v) |
|
|
|
num_heads, head_dim = self.temporal_attn.num_heads, self.temporal_attn.head_dim |
|
q = q.view(q.shape[0], num_heads, head_dim) |
|
k = k.view(k.shape[0], num_heads, head_dim) |
|
v = v.view(v.shape[0], num_heads, head_dim) |
|
|
|
cu_len = torch.cumsum(torch.tensor(split_sizes, dtype=torch.int, device=hidden_states.device), dim=0) |
|
cu_lens = [cu_len + i * N for i in range(L)] |
|
cu_lens = torch.cat([torch.zeros((1, ), device=hidden_states.device)] + cu_lens).to(torch.int) |
|
max_len = max(split_sizes) |
|
|
|
out = flash_attn_varlen_func( |
|
q=q, k=k, v=v, |
|
cu_seqlens_q=cu_lens, |
|
cu_seqlens_k=cu_lens, |
|
max_seqlen_q=max_len, |
|
max_seqlen_k=max_len, |
|
causal=self.config.temporal_causal, |
|
) |
|
|
|
out = out.view(q.shape[0], num_heads*head_dim) |
|
out = self.temporal_attn.out_proj(out) |
|
out = out.view(L, N, D).transpose(0, 1).contiguous() |
|
|
|
|
|
hidden_states = residual + out * self.alpha_xattn |
|
return hidden_states |
|
|
|
|
|
class InternVisionTemporalEncoderLayer(nn.Module): |
|
def __init__(self, config: InternVisionConfig, drop_path_rate: float, layer_idx: int=None): |
|
super().__init__() |
|
self.config = config |
|
self.layer_idx = layer_idx |
|
self.embed_dim = config.hidden_size |
|
self.intermediate_size = config.intermediate_size |
|
self.norm_type = config.norm_type |
|
|
|
self.attn = InternAttention(config) |
|
self.mlp = InternMLP(config) |
|
self.norm1 = NORM2FN[self.norm_type](self.embed_dim, eps=config.layer_norm_eps) |
|
self.norm2 = NORM2FN[self.norm_type](self.embed_dim, eps=config.layer_norm_eps) |
|
|
|
self.ls1 = nn.Parameter(config.initializer_factor * torch.ones(self.embed_dim)) |
|
self.ls2 = nn.Parameter(config.initializer_factor * torch.ones(self.embed_dim)) |
|
self.drop_path1 = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity() |
|
self.drop_path2 = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity() |
|
|
|
def initialize_temporal_module(self): |
|
temporal_layer_ids = self.config.temporal_layer_ids |
|
if (temporal_layer_ids is not None) and self.layer_idx not in temporal_layer_ids: |
|
self.temporal_module = None |
|
return |
|
|
|
self.temporal_module = TokenTemporalAttention(self.config) |
|
self.temporal_module_place = self.config.temporal_module_place |
|
param_names = [k for k, v in self.temporal_module.named_parameters()] |
|
print(f"[vision temporal model] layer {self.layer_idx} initialize temporal module. " |
|
f"Place: {self.temporal_module_place}. Parameters: {param_names}") |
|
|
|
def forward( |
|
self, |
|
hidden_states: torch.Tensor, |
|
split_sizes: Optional[list] = None, |
|
temporal_id: Optional[torch.LongTensor] = None, |
|
) -> Tuple[torch.FloatTensor, Optional[torch.FloatTensor], Optional[Tuple[torch.FloatTensor]]]: |
|
""" |
|
Args: |
|
hidden_states (`Tuple[torch.FloatTensor, Optional[torch.FloatTensor]]`): input to the layer of shape `(batch, seq_len, embed_dim)` |
|
""" |
|
if (self.temporal_module is not None) and ('before_self_attn' in self.temporal_module_place): |
|
hidden_states = self.temporal_module(hidden_states, split_sizes, temporal_id=temporal_id, place='before_self_attn') |
|
|
|
hidden_states = hidden_states + self.drop_path1(self.attn(self.norm1(hidden_states).to(hidden_states.dtype)) * self.ls1) |
|
|
|
|
|
if (self.temporal_module is not None) and ('after_self_attn' in self.temporal_module_place): |
|
hidden_states = self.temporal_module(hidden_states, split_sizes, temporal_id=temporal_id, place='after_self_attn') |
|
|
|
hidden_states = hidden_states + self.drop_path2(self.mlp(self.norm2(hidden_states).to(hidden_states.dtype)) * self.ls2) |
|
|
|
if (self.temporal_module is not None) and ('after_mlp' in self.temporal_module_place): |
|
hidden_states = self.temporal_module(hidden_states, split_sizes, temporal_id=temporal_id, place='after_mlp') |
|
|
|
return hidden_states |
|
|
|
|
|
class InternVisionTemporalEncoder(nn.Module): |
|
""" |
|
Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a |
|
[`InternEncoderLayer`]. |
|
|
|
Args: |
|
config (`InternConfig`): |
|
The corresponding vision configuration for the `InternEncoder`. |
|
""" |
|
|
|
def __init__(self, config: InternVisionConfig): |
|
super().__init__() |
|
self.config = config |
|
|
|
dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, config.num_hidden_layers)] |
|
self.layers = nn.ModuleList([ |
|
InternVisionTemporalEncoderLayer(config, dpr[idx], layer_idx=idx) |
|
for idx in range(config.num_hidden_layers) |
|
]) |
|
self.gradient_checkpointing = True |
|
|
|
def forward( |
|
self, |
|
inputs_embeds, |
|
output_hidden_states: Optional[bool] = None, |
|
return_dict: Optional[bool] = None, |
|
split_sizes: Optional[list] = None, |
|
temporal_id: Optional[torch.LongTensor] = None, |
|
) -> Union[Tuple, BaseModelOutput]: |
|
r""" |
|
Args: |
|
inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): |
|
Embedded representation of the inputs. Should be float, not int tokens. |
|
output_hidden_states (`bool`, *optional*): |
|
Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors |
|
for more detail. |
|
return_dict (`bool`, *optional*): |
|
Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. |
|
""" |
|
output_hidden_states = ( |
|
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states |
|
) |
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
|
|
|
encoder_states = () if output_hidden_states else None |
|
hidden_states = inputs_embeds |
|
|
|
for idx, encoder_layer in enumerate(self.layers): |
|
if output_hidden_states: |
|
encoder_states = encoder_states + (hidden_states,) |
|
if self.gradient_checkpointing and self.training: |
|
layer_outputs = torch.utils.checkpoint.checkpoint( |
|
encoder_layer, |
|
hidden_states, |
|
split_sizes, |
|
temporal_id) |
|
else: |
|
layer_outputs = encoder_layer( |
|
hidden_states, |
|
split_sizes=split_sizes, |
|
temporal_id=temporal_id, |
|
) |
|
hidden_states = layer_outputs |
|
|
|
if output_hidden_states: |
|
encoder_states = encoder_states + (hidden_states,) |
|
|
|
if not return_dict: |
|
return tuple(v for v in [hidden_states, encoder_states] if v is not None) |
|
return BaseModelOutput( |
|
last_hidden_state=hidden_states, hidden_states=encoder_states |
|
) |
|
|
|
|
|
class InternVisionTemporalModel(PreTrainedModel): |
|
main_input_name = 'pixel_values' |
|
_supports_flash_attn_2 = True |
|
config_class = InternVisionConfig |
|
_no_split_modules = ['InternVisionTemporalEncoderLayer'] |
|
|
|
def __init__(self, config: InternVisionConfig, delay_init_new_param=False): |
|
super().__init__(config) |
|
self.config = config |
|
|
|
self.embeddings = InternVisionEmbeddings(config) |
|
self.encoder = InternVisionTemporalEncoder(config) |
|
|
|
self.new_param_inited = False |
|
if delay_init_new_param: |
|
print(f"[vision temporal model] delay_init_new_param={delay_init_new_param}, temporal module should be initalized later") |
|
else: |
|
print(f"[vision temporal model] delay_init_new_param={delay_init_new_param}") |
|
self.initialize_temporal_module() |
|
|
|
def initialize_temporal_module(self): |
|
if self.new_param_inited: |
|
print("[vision temporal model] Warning!!! temporal modules have been initialized, skip.") |
|
return |
|
print("[vision temporal model] Initializing temporal modules...") |
|
for layer in self.encoder.layers: |
|
layer.initialize_temporal_module() |
|
self.new_param_inited = True |
|
|
|
def resize_pos_embeddings(self, old_size, new_size, patch_size): |
|
pos_emb = self.embeddings.position_embedding |
|
_, num_positions, embed_dim = pos_emb.shape |
|
cls_emb = pos_emb[:, :1, :] |
|
pos_emb = pos_emb[:, 1:, :].reshape(1, old_size // patch_size, old_size // patch_size, -1).permute(0, 3, 1, 2) |
|
pos_emb = F.interpolate(pos_emb.float(), size=new_size // patch_size, mode='bicubic', align_corners=False) |
|
pos_emb = pos_emb.to(cls_emb.dtype).reshape(1, embed_dim, -1).permute(0, 2, 1) |
|
pos_emb = torch.cat([cls_emb, pos_emb], dim=1) |
|
self.embeddings.position_embedding = nn.Parameter(pos_emb) |
|
self.embeddings.image_size = new_size |
|
logger.info('Resized position embeddings from {} to {}'.format(old_size, new_size)) |
|
|
|
def get_input_embeddings(self): |
|
return self.embeddings |
|
|
|
def forward( |
|
self, |
|
pixel_values: Optional[torch.FloatTensor] = None, |
|
output_hidden_states: Optional[bool] = None, |
|
return_dict: Optional[bool] = None, |
|
pixel_embeds: Optional[torch.FloatTensor] = None, |
|
split_sizes: Optional[list] = None, |
|
temporal_id: Optional[torch.LongTensor] = None, |
|
) -> Union[Tuple, BaseModelOutputWithPooling]: |
|
output_hidden_states = ( |
|
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states |
|
) |
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
|
|
|
if pixel_values is None and pixel_embeds is None: |
|
raise ValueError('You have to specify pixel_values or pixel_embeds') |
|
|
|
if pixel_embeds is not None: |
|
hidden_states = pixel_embeds |
|
else: |
|
if len(pixel_values.shape) == 4: |
|
hidden_states = self.embeddings(pixel_values) |
|
else: |
|
raise ValueError(f'wrong pixel_values size: {pixel_values.shape}') |
|
encoder_outputs = self.encoder( |
|
inputs_embeds=hidden_states, |
|
output_hidden_states=output_hidden_states, |
|
return_dict=return_dict, |
|
split_sizes=split_sizes, |
|
temporal_id=temporal_id, |
|
) |
|
last_hidden_state = encoder_outputs.last_hidden_state |
|
pooled_output = last_hidden_state[:, 0, :] |
|
|
|
if not return_dict: |
|
return (last_hidden_state, pooled_output) + encoder_outputs[1:] |
|
|
|
return BaseModelOutputWithPooling( |
|
last_hidden_state=last_hidden_state, |
|
pooler_output=pooled_output, |
|
hidden_states=encoder_outputs.hidden_states, |
|
attentions=encoder_outputs.attentions, |
|
) |
|
|