|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.nn.functional as F |
|
|
import numpy as np |
|
|
from functools import partial |
|
|
from dataclasses import dataclass |
|
|
from typing import Callable, Dict, Optional |
|
|
from timm.models.layers import to_2tuple |
|
|
from fairseq.tasks import FairseqTask |
|
|
from examples.data2vec.models.mae import get_2d_sincos_pos_embed, PatchEmbed |
|
|
from .base import ( |
|
|
D2vModalityConfig, |
|
|
ModalitySpecificEncoder, |
|
|
get_alibi_bias, |
|
|
MaskSeed, |
|
|
) |
|
|
from .modules import ( |
|
|
BlockEncoder, |
|
|
Decoder2d, |
|
|
FixedPositionalEncoder, |
|
|
TransformerDecoder, |
|
|
EncDecTransformerDecoder, |
|
|
) |
|
|
from examples.data2vec.data.modality import Modality |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class D2vImageConfig(D2vModalityConfig): |
|
|
type: Modality = Modality.IMAGE |
|
|
|
|
|
input_size: int = 224 |
|
|
in_chans: int = 3 |
|
|
patch_size: int = 16 |
|
|
embed_dim: int = 768 |
|
|
|
|
|
alibi_dims: int = 2 |
|
|
alibi_distance: str = "manhattan" |
|
|
|
|
|
fixed_positions: bool = True |
|
|
|
|
|
transformer_decoder: bool = False |
|
|
enc_dec_transformer: bool = False |
|
|
|
|
|
|
|
|
class ImageEncoder(ModalitySpecificEncoder): |
|
|
|
|
|
modality_cfg: D2vImageConfig |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
modality_cfg: D2vImageConfig, |
|
|
embed_dim: int, |
|
|
make_block: Callable[[float, Optional[int], Optional[int]], nn.ModuleList], |
|
|
norm_layer: Callable[[int], nn.LayerNorm], |
|
|
layer_norm_first: bool, |
|
|
alibi_biases: Dict, |
|
|
task: Optional[FairseqTask], |
|
|
): |
|
|
|
|
|
img_size = to_2tuple(modality_cfg.input_size) |
|
|
patch_size = to_2tuple(modality_cfg.patch_size) |
|
|
num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) |
|
|
|
|
|
local_encoder = PatchEmbed( |
|
|
modality_cfg.input_size, |
|
|
modality_cfg.patch_size, |
|
|
modality_cfg.in_chans, |
|
|
modality_cfg.embed_dim, |
|
|
) |
|
|
|
|
|
w = local_encoder.proj.weight.data |
|
|
torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1])) |
|
|
|
|
|
if modality_cfg.embed_dim != embed_dim: |
|
|
local_encoder = nn.Sequential( |
|
|
local_encoder, |
|
|
nn.Linear(modality_cfg.embed_dim, embed_dim), |
|
|
) |
|
|
|
|
|
project_features = nn.Identity() |
|
|
|
|
|
pos_embed = nn.Parameter( |
|
|
torch.zeros(1, num_patches, embed_dim), requires_grad=False |
|
|
) |
|
|
|
|
|
side_n = int(num_patches ** 0.5) |
|
|
|
|
|
emb = get_2d_sincos_pos_embed( |
|
|
pos_embed.shape[-1], |
|
|
side_n, |
|
|
cls_token=False, |
|
|
) |
|
|
pos_embed.data.copy_(torch.from_numpy(emb).float().unsqueeze(0)) |
|
|
fixed_positional_encoder = ( |
|
|
FixedPositionalEncoder(pos_embed) if modality_cfg.fixed_positions else None |
|
|
) |
|
|
|
|
|
dpr = np.linspace( |
|
|
modality_cfg.start_drop_path_rate, |
|
|
modality_cfg.end_drop_path_rate, |
|
|
modality_cfg.prenet_depth, |
|
|
) |
|
|
|
|
|
context_encoder = BlockEncoder( |
|
|
nn.ModuleList(make_block(dpr[i]) for i in range(modality_cfg.prenet_depth)), |
|
|
norm_layer(embed_dim) if not layer_norm_first else None, |
|
|
layer_norm_first, |
|
|
modality_cfg.prenet_layerdrop, |
|
|
modality_cfg.prenet_dropout, |
|
|
) |
|
|
|
|
|
if modality_cfg.transformer_decoder: |
|
|
if modality_cfg.enc_dec_transformer: |
|
|
decoder = EncDecTransformerDecoder(modality_cfg.decoder, embed_dim) |
|
|
else: |
|
|
dec_enc = BlockEncoder( |
|
|
nn.ModuleList( |
|
|
make_block(0, modality_cfg.decoder.decoder_dim, 8) |
|
|
for _ in range(modality_cfg.decoder.decoder_layers) |
|
|
), |
|
|
None, |
|
|
layer_norm_first, |
|
|
0, |
|
|
0, |
|
|
) |
|
|
decoder = TransformerDecoder(modality_cfg.decoder, embed_dim, dec_enc) |
|
|
else: |
|
|
decoder = ( |
|
|
Decoder2d(modality_cfg.decoder, embed_dim, side_n, side_n) |
|
|
if modality_cfg.decoder is not None |
|
|
else None |
|
|
) |
|
|
|
|
|
alibi_bias_fn = partial( |
|
|
get_alibi_bias, |
|
|
alibi_biases=alibi_biases, |
|
|
heads=modality_cfg.num_alibi_heads, |
|
|
dims=modality_cfg.alibi_dims, |
|
|
distance=modality_cfg.alibi_distance, |
|
|
) |
|
|
|
|
|
super().__init__( |
|
|
modality_cfg=modality_cfg, |
|
|
embed_dim=embed_dim, |
|
|
local_encoder=local_encoder, |
|
|
project_features=project_features, |
|
|
fixed_positional_encoder=fixed_positional_encoder, |
|
|
relative_positional_encoder=None, |
|
|
context_encoder=context_encoder, |
|
|
decoder=decoder, |
|
|
get_alibi_bias=alibi_bias_fn, |
|
|
) |
|
|
|
|
|
def reset_parameters(self): |
|
|
super().reset_parameters() |
|
|
if self.decoder is not None: |
|
|
self.decoder.reset_parameters() |
|
|
|
|
|
@torch.no_grad() |
|
|
def patchify(self, imgs): |
|
|
""" |
|
|
imgs: (N, 3, H, W) |
|
|
x: (N, L, patch_size**2 *3) |
|
|
""" |
|
|
p = self.modality_cfg.patch_size |
|
|
h = w = imgs.shape[2] // p |
|
|
x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p)) |
|
|
x = torch.einsum("nchpwq->nhwpqc", x) |
|
|
x = x.reshape(shape=(imgs.shape[0], h * w, p ** 2 * 3)) |
|
|
|
|
|
return x |
|
|
|
|
|
@torch.no_grad() |
|
|
def unpatchify(self, x): |
|
|
""" |
|
|
x: (N, L, patch_size**2 *3) |
|
|
imgs: (N, 3, H, W) |
|
|
""" |
|
|
p = self.modality_cfg.patch_size |
|
|
h = w = int(x.shape[1] ** 0.5) |
|
|
assert h * w == x.shape[1] |
|
|
|
|
|
x = x.reshape(shape=(x.shape[0], h, w, p, p, 3)) |
|
|
x = torch.einsum("nhwpqc->nchpwq", x) |
|
|
imgs = x.reshape(shape=(x.shape[0], 3, h * p, h * p)) |
|
|
return imgs |
|
|
|
|
|
def compute_mask( |
|
|
self, |
|
|
x, |
|
|
padding_mask, |
|
|
mask_seed: Optional[MaskSeed], |
|
|
apply, |
|
|
shape=None, |
|
|
precomputed_mask=None, |
|
|
): |
|
|
mlen = self.modality_cfg.mask_length |
|
|
if mlen <= 1: |
|
|
return super().compute_mask( |
|
|
x, padding_mask, mask_seed, apply, precomputed_mask |
|
|
) |
|
|
|
|
|
if precomputed_mask is not None: |
|
|
mask = precomputed_mask |
|
|
else: |
|
|
from fairseq.data.data_utils import compute_block_mask_2d |
|
|
|
|
|
if shape is not None: |
|
|
B, L, D = shape |
|
|
else: |
|
|
B, L, D = x.shape |
|
|
|
|
|
mask = compute_block_mask_2d( |
|
|
shape=(B, L), |
|
|
mask_prob=self.modality_cfg.mask_prob, |
|
|
mask_length=self.modality_cfg.mask_length, |
|
|
mask_prob_adjust=self.modality_cfg.mask_prob_adjust, |
|
|
inverse_mask=self.modality_cfg.inverse_mask, |
|
|
require_same_masks=True, |
|
|
mask_dropout=self.modality_cfg.mask_dropout, |
|
|
) |
|
|
|
|
|
mask_info = self.make_maskinfo(x, mask, shape) |
|
|
if apply: |
|
|
x = self.apply_mask(x, mask_info) |
|
|
|
|
|
return x, mask_info |
|
|
|
|
|
def decoder_input(self, x, mask_info): |
|
|
if ( |
|
|
not self.modality_cfg.transformer_decoder |
|
|
or not self.modality_cfg.enc_dec_transformer |
|
|
): |
|
|
return super().decoder_input(x, mask_info) |
|
|
|
|
|
inp_drop = self.modality_cfg.decoder.input_dropout |
|
|
if inp_drop > 0: |
|
|
x = F.dropout(x, inp_drop, training=self.training, inplace=True) |
|
|
|
|
|
kv = x[:, self.modality_cfg.num_extra_tokens :] |
|
|
|
|
|
assert self.fixed_positional_encoder is not None |
|
|
pos = self.fixed_positional_encoder(x, None).expand(x.size(0), -1, -1) |
|
|
|
|
|
mask = mask_info.mask.bool() |
|
|
if self.modality_cfg.decoder.add_positions_all: |
|
|
kv = kv + pos[~mask].view(kv.shape) |
|
|
|
|
|
q = pos[mask].view(x.size(0), -1, x.size(-1)) |
|
|
|
|
|
return q, kv |
|
|
|
