comi-logic/Wan2GP
7
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

stuff and more stuff

Browse Source
This commit is contained in:
DeepBeepMeep
2025-07-02 14:08:59 +02:00
parent e0666a3e6d
commit 28fc48db2d
125 changed files with 11593 additions and 578 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
postprocessing/mmaudio/ext/synchformer/LICENSE Normal file
View File

@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2024 Vladimir Iashin
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

1
postprocessing/mmaudio/ext/synchformer/__init__.py Normal file
View File

@@ -0,0 +1 @@
# from .synchformer import Synchformer

84
postprocessing/mmaudio/ext/synchformer/divided_224_16x4.yaml Normal file
View File

@@ -0,0 +1,84 @@
TRAIN:
ENABLE: True
DATASET: Ssv2
BATCH_SIZE: 32
EVAL_PERIOD: 5
CHECKPOINT_PERIOD: 5
AUTO_RESUME: True
CHECKPOINT_EPOCH_RESET: True
CHECKPOINT_FILE_PATH: /checkpoint/fmetze/neurips_sota/40944587/checkpoints/checkpoint_epoch_00035.pyth
DATA:
NUM_FRAMES: 16
SAMPLING_RATE: 4
TRAIN_JITTER_SCALES: [256, 320]
TRAIN_CROP_SIZE: 224
TEST_CROP_SIZE: 224
INPUT_CHANNEL_NUM: [3]
MEAN: [0.5, 0.5, 0.5]
STD: [0.5, 0.5, 0.5]
PATH_TO_DATA_DIR: /private/home/mandelapatrick/slowfast/data/ssv2
PATH_PREFIX: /datasets01/SomethingV2/092720/20bn-something-something-v2-frames
INV_UNIFORM_SAMPLE: True
RANDOM_FLIP: False
REVERSE_INPUT_CHANNEL: True
USE_RAND_AUGMENT: True
RE_PROB: 0.0
USE_REPEATED_AUG: False
USE_RANDOM_RESIZE_CROPS: False
COLORJITTER: False
GRAYSCALE: False
GAUSSIAN: False
SOLVER:
BASE_LR: 1e-4
LR_POLICY: steps_with_relative_lrs
LRS: [1, 0.1, 0.01]
STEPS: [0, 20, 30]
MAX_EPOCH: 35
MOMENTUM: 0.9
WEIGHT_DECAY: 5e-2
WARMUP_EPOCHS: 0.0
OPTIMIZING_METHOD: adamw
USE_MIXED_PRECISION: True
SMOOTHING: 0.2
SLOWFAST:
ALPHA: 8
VIT:
PATCH_SIZE: 16
PATCH_SIZE_TEMP: 2
CHANNELS: 3
EMBED_DIM: 768
DEPTH: 12
NUM_HEADS: 12
MLP_RATIO: 4
QKV_BIAS: True
VIDEO_INPUT: True
TEMPORAL_RESOLUTION: 8
USE_MLP: True
DROP: 0.0
POS_DROPOUT: 0.0
DROP_PATH: 0.2
IM_PRETRAINED: True
HEAD_DROPOUT: 0.0
HEAD_ACT: tanh
PRETRAINED_WEIGHTS: vit_1k
ATTN_LAYER: divided
MODEL:
NUM_CLASSES: 174
ARCH: slow
MODEL_NAME: VisionTransformer
LOSS_FUNC: cross_entropy
TEST:
ENABLE: True
DATASET: Ssv2
BATCH_SIZE: 64
NUM_ENSEMBLE_VIEWS: 1
NUM_SPATIAL_CROPS: 3
DATA_LOADER:
NUM_WORKERS: 4
PIN_MEMORY: True
NUM_GPUS: 8
NUM_SHARDS: 4
RNG_SEED: 0
OUTPUT_DIR: .
TENSORBOARD:
ENABLE: True

400
postprocessing/mmaudio/ext/synchformer/motionformer.py Normal file
View File

@@ -0,0 +1,400 @@
import logging
from pathlib import Path
import einops
import torch
from omegaconf import OmegaConf
from timm.layers import trunc_normal_
from torch import nn
from .utils import check_if_file_exists_else_download
from .video_model_builder import VisionTransformer
FILE2URL = {
# cfg
'motionformer_224_16x4.yaml':
'https://raw.githubusercontent.com/facebookresearch/Motionformer/bf43d50/configs/SSV2/motionformer_224_16x4.yaml',
'joint_224_16x4.yaml':
'https://raw.githubusercontent.com/facebookresearch/Motionformer/bf43d50/configs/SSV2/joint_224_16x4.yaml',
'divided_224_16x4.yaml':
'https://raw.githubusercontent.com/facebookresearch/Motionformer/bf43d50/configs/SSV2/divided_224_16x4.yaml',
# ckpt
'ssv2_motionformer_224_16x4.pyth':
'https://dl.fbaipublicfiles.com/motionformer/ssv2_motionformer_224_16x4.pyth',
'ssv2_joint_224_16x4.pyth':
'https://dl.fbaipublicfiles.com/motionformer/ssv2_joint_224_16x4.pyth',
'ssv2_divided_224_16x4.pyth':
'https://dl.fbaipublicfiles.com/motionformer/ssv2_divided_224_16x4.pyth',
}
class MotionFormer(VisionTransformer):
''' This class serves three puposes:
1. Renames the class to MotionFormer.
2. Downloads the cfg from the original repo and patches it if needed.
3. Takes care of feature extraction by redefining .forward()
- if `extract_features=True` and `factorize_space_time=False`,
the output is of shape (B, T, D) where T = 1 + (224 // 16) * (224 // 16) * 8
- if `extract_features=True` and `factorize_space_time=True`, the output is of shape (B*S, D)
and spatial and temporal transformer encoder layers are used.
- if `extract_features=True` and `factorize_space_time=True` as well as `add_global_repr=True`
the output is of shape (B, D) and spatial and temporal transformer encoder layers
are used as well as the global representation is extracted from segments (extra pos emb
is added).
'''
def __init__(
self,
extract_features: bool = False,
ckpt_path: str = None,
factorize_space_time: bool = None,
agg_space_module: str = None,
agg_time_module: str = None,
add_global_repr: bool = True,
agg_segments_module: str = None,
max_segments: int = None,
):
self.extract_features = extract_features
self.ckpt_path = ckpt_path
self.factorize_space_time = factorize_space_time
if self.ckpt_path is not None:
check_if_file_exists_else_download(self.ckpt_path, FILE2URL)
ckpt = torch.load(self.ckpt_path, map_location='cpu')
mformer_ckpt2cfg = {
'ssv2_motionformer_224_16x4.pyth': 'motionformer_224_16x4.yaml',
'ssv2_joint_224_16x4.pyth': 'joint_224_16x4.yaml',
'ssv2_divided_224_16x4.pyth': 'divided_224_16x4.yaml',
}
# init from motionformer ckpt or from our Stage I ckpt
# depending on whether the feat extractor was pre-trained on AVCLIPMoCo or not, we need to
# load the state dict differently
was_pt_on_avclip = self.ckpt_path.endswith(
'.pt') # checks if it is a stage I ckpt (FIXME: a bit generic)
if self.ckpt_path.endswith(tuple(mformer_ckpt2cfg.keys())):
cfg_fname = mformer_ckpt2cfg[Path(self.ckpt_path).name]
elif was_pt_on_avclip:
# TODO: this is a hack, we should be able to get the cfg from the ckpt (earlier ckpt didn't have it)
s1_cfg = ckpt.get('args', None) # Stage I cfg
if s1_cfg is not None:
s1_vfeat_extractor_ckpt_path = s1_cfg.model.params.vfeat_extractor.params.ckpt_path
# if the stage I ckpt was initialized from a motionformer ckpt or train from scratch
if s1_vfeat_extractor_ckpt_path is not None:
cfg_fname = mformer_ckpt2cfg[Path(s1_vfeat_extractor_ckpt_path).name]
else:
cfg_fname = 'divided_224_16x4.yaml'
else:
cfg_fname = 'divided_224_16x4.yaml'
else:
raise ValueError(f'ckpt_path {self.ckpt_path} is not supported.')
else:
was_pt_on_avclip = False
cfg_fname = 'divided_224_16x4.yaml'
# logging.info(f'No ckpt_path provided, using {cfg_fname} config.')
if cfg_fname in ['motionformer_224_16x4.yaml', 'divided_224_16x4.yaml']:
pos_emb_type = 'separate'
elif cfg_fname == 'joint_224_16x4.yaml':
pos_emb_type = 'joint'
self.mformer_cfg_path = Path(__file__).absolute().parent / cfg_fname
check_if_file_exists_else_download(self.mformer_cfg_path, FILE2URL)
mformer_cfg = OmegaConf.load(self.mformer_cfg_path)
logging.info(f'Loading MotionFormer config from {self.mformer_cfg_path.absolute()}')
# patch the cfg (from the default cfg defined in the repo `Motionformer/slowfast/config/defaults.py`)
mformer_cfg.VIT.ATTN_DROPOUT = 0.0
mformer_cfg.VIT.POS_EMBED = pos_emb_type
mformer_cfg.VIT.USE_ORIGINAL_TRAJ_ATTN_CODE = True
mformer_cfg.VIT.APPROX_ATTN_TYPE = 'none' # guessing
mformer_cfg.VIT.APPROX_ATTN_DIM = 64 # from ckpt['cfg']
# finally init VisionTransformer with the cfg
super().__init__(mformer_cfg)
# load the ckpt now if ckpt is provided and not from AVCLIPMoCo-pretrained ckpt
if (self.ckpt_path is not None) and (not was_pt_on_avclip):
_ckpt_load_status = self.load_state_dict(ckpt['model_state'], strict=False)
if len(_ckpt_load_status.missing_keys) > 0 or len(
_ckpt_load_status.unexpected_keys) > 0:
logging.warning(f'Loading exact vfeat_extractor ckpt from {self.ckpt_path} failed.' \
f'Missing keys: {_ckpt_load_status.missing_keys}, ' \
f'Unexpected keys: {_ckpt_load_status.unexpected_keys}')
else:
logging.info(f'Loading vfeat_extractor ckpt from {self.ckpt_path} succeeded.')
if self.extract_features:
assert isinstance(self.norm,
nn.LayerNorm), 'early x[:, 1:, :] may not be safe for per-tr weights'
# pre-logits are Sequential(nn.Linear(emb, emd), act) and `act` is tanh but see the logger
self.pre_logits = nn.Identity()
# we don't need the classification head (saving memory)
self.head = nn.Identity()
self.head_drop = nn.Identity()
# avoiding code duplication (used only if agg_*_module is TransformerEncoderLayer)
transf_enc_layer_kwargs = dict(
d_model=self.embed_dim,
nhead=self.num_heads,
activation=nn.GELU(),
batch_first=True,
dim_feedforward=self.mlp_ratio * self.embed_dim,
dropout=self.drop_rate,
layer_norm_eps=1e-6,
norm_first=True,
)
# define adapters if needed
if self.factorize_space_time:
if agg_space_module == 'TransformerEncoderLayer':
self.spatial_attn_agg = SpatialTransformerEncoderLayer(
**transf_enc_layer_kwargs)
elif agg_space_module == 'AveragePooling':
self.spatial_attn_agg = AveragePooling(avg_pattern='BS D t h w -> BS D t',
then_permute_pattern='BS D t -> BS t D')
if agg_time_module == 'TransformerEncoderLayer':
self.temp_attn_agg = TemporalTransformerEncoderLayer(**transf_enc_layer_kwargs)
elif agg_time_module == 'AveragePooling':
self.temp_attn_agg = AveragePooling(avg_pattern='BS t D -> BS D')
elif 'Identity' in agg_time_module:
self.temp_attn_agg = nn.Identity()
# define a global aggregation layer (aggregarate over segments)
self.add_global_repr = add_global_repr
if add_global_repr:
if agg_segments_module == 'TransformerEncoderLayer':
# we can reuse the same layer as for temporal factorization (B, dim_to_agg, D) -> (B, D)
# we need to add pos emb (PE) because previously we added the same PE for each segment
pos_max_len = max_segments if max_segments is not None else 16 # 16 = 10sec//0.64sec + 1
self.global_attn_agg = TemporalTransformerEncoderLayer(
add_pos_emb=True,
pos_emb_drop=mformer_cfg.VIT.POS_DROPOUT,
pos_max_len=pos_max_len,
**transf_enc_layer_kwargs)
elif agg_segments_module == 'AveragePooling':
self.global_attn_agg = AveragePooling(avg_pattern='B S D -> B D')
if was_pt_on_avclip:
# we need to filter out the state_dict of the AVCLIP model (has both A and V extractors)
# and keep only the state_dict of the feat extractor
ckpt_weights = dict()
for k, v in ckpt['state_dict'].items():
if k.startswith(('module.v_encoder.', 'v_encoder.')):
k = k.replace('module.', '').replace('v_encoder.', '')
ckpt_weights[k] = v
_load_status = self.load_state_dict(ckpt_weights, strict=False)
if len(_load_status.missing_keys) > 0 or len(_load_status.unexpected_keys) > 0:
logging.warning(f'Loading exact vfeat_extractor ckpt from {self.ckpt_path} failed. \n' \
f'Missing keys ({len(_load_status.missing_keys)}): ' \
f'{_load_status.missing_keys}, \n' \
f'Unexpected keys ({len(_load_status.unexpected_keys)}): ' \
f'{_load_status.unexpected_keys} \n' \
f'temp_attn_agg are expected to be missing if ckpt was pt contrastively.')
else:
logging.info(f'Loading vfeat_extractor ckpt from {self.ckpt_path} succeeded.')
# patch_embed is not used in MotionFormer, only patch_embed_3d, because cfg.VIT.PATCH_SIZE_TEMP > 1
# but it used to calculate the number of patches, so we need to set keep it
self.patch_embed.requires_grad_(False)
def forward(self, x):
'''
x is of shape (B, S, C, T, H, W) where S is the number of segments.
'''
# Batch, Segments, Channels, T=frames, Height, Width
B, S, C, T, H, W = x.shape
# Motionformer expects a tensor of shape (1, B, C, T, H, W).
# The first dimension (1) is a dummy dimension to make the input tensor and won't be used:
# see `video_model_builder.video_input`.
# x = x.unsqueeze(0) # (1, B, S, C, T, H, W)
orig_shape = (B, S, C, T, H, W)
x = x.view(B * S, C, T, H, W) # flatten batch and segments
x = self.forward_segments(x, orig_shape=orig_shape)
# unpack the segments (using rest dimensions to support different shapes e.g. (BS, D) or (BS, t, D))
x = x.view(B, S, *x.shape[1:])
# x is now of shape (B*S, D) or (B*S, t, D) if `self.temp_attn_agg` is `Identity`
return x # x is (B, S, ...)
def forward_segments(self, x, orig_shape: tuple) -> torch.Tensor:
'''x is of shape (1, BS, C, T, H, W) where S is the number of segments.'''
x, x_mask = self.forward_features(x)
assert self.extract_features
# (BS, T, D) where T = 1 + (224 // 16) * (224 // 16) * 8
x = x[:,
1:, :] # without the CLS token for efficiency (should be safe for LayerNorm and FC)
x = self.norm(x)
x = self.pre_logits(x)
if self.factorize_space_time:
x = self.restore_spatio_temp_dims(x, orig_shape) # (B*S, D, t, h, w) <- (B*S, t*h*w, D)
x = self.spatial_attn_agg(x, x_mask) # (B*S, t, D)
x = self.temp_attn_agg(
x) # (B*S, D) or (BS, t, D) if `self.temp_attn_agg` is `Identity`
return x
def restore_spatio_temp_dims(self, feats: torch.Tensor, orig_shape: tuple) -> torch.Tensor:
'''
feats are of shape (B*S, T, D) where T = 1 + (224 // 16) * (224 // 16) * 8
Our goal is to make them of shape (B*S, t, h, w, D) where h, w are the spatial dimensions.
From `self.patch_embed_3d`, it follows that we could reshape feats with:
`feats.transpose(1, 2).view(B*S, D, t, h, w)`
'''
B, S, C, T, H, W = orig_shape
D = self.embed_dim
# num patches in each dimension
t = T // self.patch_embed_3d.z_block_size
h = self.patch_embed_3d.height
w = self.patch_embed_3d.width
feats = feats.permute(0, 2, 1) # (B*S, D, T)
feats = feats.view(B * S, D, t, h, w) # (B*S, D, t, h, w)
return feats
class BaseEncoderLayer(nn.TransformerEncoderLayer):
'''
This is a wrapper around nn.TransformerEncoderLayer that adds a CLS token
to the sequence and outputs the CLS token's representation.
This base class parents both SpatialEncoderLayer and TemporalEncoderLayer for the RGB stream
and the FrequencyEncoderLayer and TemporalEncoderLayer for the audio stream stream.
We also, optionally, add a positional embedding to the input sequence which
allows to reuse it for global aggregation (of segments) for both streams.
'''
def __init__(self,
add_pos_emb: bool = False,
pos_emb_drop: float = None,
pos_max_len: int = None,
*args_transformer_enc,
**kwargs_transformer_enc):
super().__init__(*args_transformer_enc, **kwargs_transformer_enc)
self.cls_token = nn.Parameter(torch.zeros(1, 1, self.self_attn.embed_dim))
trunc_normal_(self.cls_token, std=.02)
# add positional embedding
self.add_pos_emb = add_pos_emb
if add_pos_emb:
self.pos_max_len = 1 + pos_max_len # +1 (for CLS)
self.pos_emb = nn.Parameter(torch.zeros(1, self.pos_max_len, self.self_attn.embed_dim))
self.pos_drop = nn.Dropout(pos_emb_drop)
trunc_normal_(self.pos_emb, std=.02)
self.apply(self._init_weights)
def forward(self, x: torch.Tensor, x_mask: torch.Tensor = None):
''' x is of shape (B, N, D); if provided x_mask is of shape (B, N)'''
batch_dim = x.shape[0]
# add CLS token
cls_tokens = self.cls_token.expand(batch_dim, -1, -1) # expanding to match batch dimension
x = torch.cat((cls_tokens, x), dim=-2) # (batch_dim, 1+seq_len, D)
if x_mask is not None:
cls_mask = torch.ones((batch_dim, 1), dtype=torch.bool,
device=x_mask.device) # 1=keep; 0=mask
x_mask_w_cls = torch.cat((cls_mask, x_mask), dim=-1) # (batch_dim, 1+seq_len)
B, N = x_mask_w_cls.shape
# torch expects (N, N) or (B*num_heads, N, N) mask (sadness ahead); torch masks
x_mask_w_cls = x_mask_w_cls.reshape(B, 1, 1, N)\
.expand(-1, self.self_attn.num_heads, N, -1)\
.reshape(B * self.self_attn.num_heads, N, N)
assert x_mask_w_cls.dtype == x_mask_w_cls.bool().dtype, 'x_mask_w_cls.dtype != bool'
x_mask_w_cls = ~x_mask_w_cls # invert mask (1=mask)
else:
x_mask_w_cls = None
# add positional embedding
if self.add_pos_emb:
seq_len = x.shape[
1] # (don't even think about moving it before the CLS token concatenation)
assert seq_len <= self.pos_max_len, f'Seq len ({seq_len}) > pos_max_len ({self.pos_max_len})'
x = x + self.pos_emb[:, :seq_len, :]
x = self.pos_drop(x)
# apply encoder layer (calls nn.TransformerEncoderLayer.forward);
x = super().forward(src=x, src_mask=x_mask_w_cls) # (batch_dim, 1+seq_len, D)
# CLS token is expected to hold spatial information for each frame
x = x[:, 0, :] # (batch_dim, D)
return x
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
@torch.jit.ignore
def no_weight_decay(self):
return {'cls_token', 'pos_emb'}
class SpatialTransformerEncoderLayer(BaseEncoderLayer):
''' Aggregates spatial dimensions by applying attention individually to each frame. '''
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def forward(self, x: torch.Tensor, x_mask: torch.Tensor = None) -> torch.Tensor:
''' x is of shape (B*S, D, t, h, w) where S is the number of segments.
if specified x_mask (B*S, t, h, w), 0=masked, 1=kept
Returns a tensor of shape (B*S, t, D) pooling spatial information for each frame. '''
BS, D, t, h, w = x.shape
# time as a batch dimension and flatten spatial dimensions as sequence
x = einops.rearrange(x, 'BS D t h w -> (BS t) (h w) D')
# similar to mask
if x_mask is not None:
x_mask = einops.rearrange(x_mask, 'BS t h w -> (BS t) (h w)')
# apply encoder layer (BaseEncoderLayer.forward) - it will add CLS token and output its representation
x = super().forward(x=x, x_mask=x_mask) # (B*S*t, D)
# reshape back to (B*S, t, D)
x = einops.rearrange(x, '(BS t) D -> BS t D', BS=BS, t=t)
# (B*S, t, D)
return x
class TemporalTransformerEncoderLayer(BaseEncoderLayer):
''' Aggregates temporal dimension with attention. Also used with pos emb as global aggregation
in both streams. '''
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def forward(self, x):
''' x is of shape (B*S, t, D) where S is the number of segments.
Returns a tensor of shape (B*S, D) pooling temporal information. '''
BS, t, D = x.shape
# apply encoder layer (BaseEncoderLayer.forward) - it will add CLS token and output its representation
x = super().forward(x) # (B*S, D)
return x # (B*S, D)
class AveragePooling(nn.Module):
def __init__(self, avg_pattern: str, then_permute_pattern: str = None) -> None:
''' patterns are e.g. "bs t d -> bs d" '''
super().__init__()
# TODO: need to register them as buffers (but fails because these are strings)
self.reduce_fn = 'mean'
self.avg_pattern = avg_pattern
self.then_permute_pattern = then_permute_pattern
def forward(self, x: torch.Tensor, x_mask: torch.Tensor = None) -> torch.Tensor:
x = einops.reduce(x, self.avg_pattern, self.reduce_fn)
if self.then_permute_pattern is not None:
x = einops.rearrange(x, self.then_permute_pattern)
return x

55
postprocessing/mmaudio/ext/synchformer/synchformer.py Normal file
View File

@@ -0,0 +1,55 @@
import logging
from typing import Any, Mapping
import torch
from torch import nn
from .motionformer import MotionFormer
class Synchformer(nn.Module):
def __init__(self):
super().__init__()
self.vfeat_extractor = MotionFormer(extract_features=True,
factorize_space_time=True,
agg_space_module='TransformerEncoderLayer',
agg_time_module='torch.nn.Identity',
add_global_repr=False)
# self.vfeat_extractor = instantiate_from_config(vfeat_extractor)
# self.afeat_extractor = instantiate_from_config(afeat_extractor)
# # bridging the s3d latent dim (1024) into what is specified in the config
# # to match e.g. the transformer dim
# self.vproj = instantiate_from_config(vproj)
# self.aproj = instantiate_from_config(aproj)
# self.transformer = instantiate_from_config(transformer)
def forward(self, vis):
B, S, Tv, C, H, W = vis.shape
vis = vis.permute(0, 1, 3, 2, 4, 5) # (B, S, C, Tv, H, W)
# feat extractors return a tuple of segment-level and global features (ignored for sync)
# (B, S, tv, D), e.g. (B, 7, 8, 768)
vis = self.vfeat_extractor(vis)
return vis
def load_state_dict(self, sd: Mapping[str, Any], strict: bool = True):
# discard all entries except vfeat_extractor
sd = {k: v for k, v in sd.items() if k.startswith('vfeat_extractor')}
return super().load_state_dict(sd, strict)
if __name__ == "__main__":
model = Synchformer().cuda().eval()
sd = torch.load('./ext_weights/synchformer_state_dict.pth', weights_only=True)
model.load_state_dict(sd)
vid = torch.randn(2, 7, 16, 3, 224, 224).cuda()
features = model.extract_vfeats(vid, for_loop=False).detach().cpu()
print(features.shape)
# extract and save the state dict only
# sd = torch.load('./ext_weights/sync_model_audioset.pt')['model']
# torch.save(sd, './ext_weights/synchformer_state_dict.pth')

92
postprocessing/mmaudio/ext/synchformer/utils.py Normal file
View File

@@ -0,0 +1,92 @@
from hashlib import md5
from pathlib import Path
import requests
from tqdm import tqdm
PARENT_LINK = 'https://a3s.fi/swift/v1/AUTH_a235c0f452d648828f745589cde1219a'
FNAME2LINK = {
# S3: Synchability: AudioSet (run 2)
'24-01-22T20-34-52.pt':
f'{PARENT_LINK}/sync/sync_models/24-01-22T20-34-52/24-01-22T20-34-52.pt',
'cfg-24-01-22T20-34-52.yaml':
f'{PARENT_LINK}/sync/sync_models/24-01-22T20-34-52/cfg-24-01-22T20-34-52.yaml',
# S2: Synchformer: AudioSet (run 2)
'24-01-04T16-39-21.pt':
f'{PARENT_LINK}/sync/sync_models/24-01-04T16-39-21/24-01-04T16-39-21.pt',
'cfg-24-01-04T16-39-21.yaml':
f'{PARENT_LINK}/sync/sync_models/24-01-04T16-39-21/cfg-24-01-04T16-39-21.yaml',
# S2: Synchformer: AudioSet (run 1)
'23-08-28T11-23-23.pt':
f'{PARENT_LINK}/sync/sync_models/23-08-28T11-23-23/23-08-28T11-23-23.pt',
'cfg-23-08-28T11-23-23.yaml':
f'{PARENT_LINK}/sync/sync_models/23-08-28T11-23-23/cfg-23-08-28T11-23-23.yaml',
# S2: Synchformer: LRS3 (run 2)
'23-12-23T18-33-57.pt':
f'{PARENT_LINK}/sync/sync_models/23-12-23T18-33-57/23-12-23T18-33-57.pt',
'cfg-23-12-23T18-33-57.yaml':
f'{PARENT_LINK}/sync/sync_models/23-12-23T18-33-57/cfg-23-12-23T18-33-57.yaml',
# S2: Synchformer: VGS (run 2)
'24-01-02T10-00-53.pt':
f'{PARENT_LINK}/sync/sync_models/24-01-02T10-00-53/24-01-02T10-00-53.pt',
'cfg-24-01-02T10-00-53.yaml':
f'{PARENT_LINK}/sync/sync_models/24-01-02T10-00-53/cfg-24-01-02T10-00-53.yaml',
# SparseSync: ft VGGSound-Full
'22-09-21T21-00-52.pt':
f'{PARENT_LINK}/sync/sync_models/22-09-21T21-00-52/22-09-21T21-00-52.pt',
'cfg-22-09-21T21-00-52.yaml':
f'{PARENT_LINK}/sync/sync_models/22-09-21T21-00-52/cfg-22-09-21T21-00-52.yaml',
# SparseSync: ft VGGSound-Sparse
'22-07-28T15-49-45.pt':
f'{PARENT_LINK}/sync/sync_models/22-07-28T15-49-45/22-07-28T15-49-45.pt',
'cfg-22-07-28T15-49-45.yaml':
f'{PARENT_LINK}/sync/sync_models/22-07-28T15-49-45/cfg-22-07-28T15-49-45.yaml',
# SparseSync: only pt on LRS3
'22-07-13T22-25-49.pt':
f'{PARENT_LINK}/sync/sync_models/22-07-13T22-25-49/22-07-13T22-25-49.pt',
'cfg-22-07-13T22-25-49.yaml':
f'{PARENT_LINK}/sync/sync_models/22-07-13T22-25-49/cfg-22-07-13T22-25-49.yaml',
# SparseSync: feature extractors
'ResNetAudio-22-08-04T09-51-04.pt':
f'{PARENT_LINK}/sync/ResNetAudio-22-08-04T09-51-04.pt', # 2s
'ResNetAudio-22-08-03T23-14-49.pt':
f'{PARENT_LINK}/sync/ResNetAudio-22-08-03T23-14-49.pt', # 3s
'ResNetAudio-22-08-03T23-14-28.pt':
f'{PARENT_LINK}/sync/ResNetAudio-22-08-03T23-14-28.pt', # 4s
'ResNetAudio-22-06-24T08-10-33.pt':
f'{PARENT_LINK}/sync/ResNetAudio-22-06-24T08-10-33.pt', # 5s
'ResNetAudio-22-06-24T17-31-07.pt':
f'{PARENT_LINK}/sync/ResNetAudio-22-06-24T17-31-07.pt', # 6s
'ResNetAudio-22-06-24T23-57-11.pt':
f'{PARENT_LINK}/sync/ResNetAudio-22-06-24T23-57-11.pt', # 7s
'ResNetAudio-22-06-25T04-35-42.pt':
f'{PARENT_LINK}/sync/ResNetAudio-22-06-25T04-35-42.pt', # 8s
}
def check_if_file_exists_else_download(path, fname2link=FNAME2LINK, chunk_size=1024):
'''Checks if file exists, if not downloads it from the link to the path'''
path = Path(path)
if not path.exists():
path.parent.mkdir(exist_ok=True, parents=True)
link = fname2link.get(path.name, None)
if link is None:
raise ValueError(f'Cant find the checkpoint file: {path}.',
f'Please download it manually and ensure the path exists.')
with requests.get(fname2link[path.name], stream=True) as r:
total_size = int(r.headers.get('content-length', 0))
with tqdm(total=total_size, unit='B', unit_scale=True) as pbar:
with open(path, 'wb') as f:
for data in r.iter_content(chunk_size=chunk_size):
if data:
f.write(data)
pbar.update(chunk_size)
def get_md5sum(path):
hash_md5 = md5()
with open(path, 'rb') as f:
for chunk in iter(lambda: f.read(4096 * 8), b''):
hash_md5.update(chunk)
md5sum = hash_md5.hexdigest()
return md5sum

277
postprocessing/mmaudio/ext/synchformer/video_model_builder.py Normal file
View File

@@ -0,0 +1,277 @@
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
# Copyright 2020 Ross Wightman
# Modified Model definition
from collections import OrderedDict
from functools import partial
import torch
import torch.nn as nn
from timm.layers import trunc_normal_
from . import vit_helper
class VisionTransformer(nn.Module):
""" Vision Transformer with support for patch or hybrid CNN input stage """
def __init__(self, cfg):
super().__init__()
self.img_size = cfg.DATA.TRAIN_CROP_SIZE
self.patch_size = cfg.VIT.PATCH_SIZE
self.in_chans = cfg.VIT.CHANNELS
if cfg.TRAIN.DATASET == "Epickitchens":
self.num_classes = [97, 300]
else:
self.num_classes = cfg.MODEL.NUM_CLASSES
self.embed_dim = cfg.VIT.EMBED_DIM
self.depth = cfg.VIT.DEPTH
self.num_heads = cfg.VIT.NUM_HEADS
self.mlp_ratio = cfg.VIT.MLP_RATIO
self.qkv_bias = cfg.VIT.QKV_BIAS
self.drop_rate = cfg.VIT.DROP
self.drop_path_rate = cfg.VIT.DROP_PATH
self.head_dropout = cfg.VIT.HEAD_DROPOUT
self.video_input = cfg.VIT.VIDEO_INPUT
self.temporal_resolution = cfg.VIT.TEMPORAL_RESOLUTION
self.use_mlp = cfg.VIT.USE_MLP
self.num_features = self.embed_dim
norm_layer = partial(nn.LayerNorm, eps=1e-6)
self.attn_drop_rate = cfg.VIT.ATTN_DROPOUT
self.head_act = cfg.VIT.HEAD_ACT
self.cfg = cfg
# Patch Embedding
self.patch_embed = vit_helper.PatchEmbed(img_size=224,
patch_size=self.patch_size,
in_chans=self.in_chans,
embed_dim=self.embed_dim)
# 3D Patch Embedding
self.patch_embed_3d = vit_helper.PatchEmbed3D(img_size=self.img_size,
temporal_resolution=self.temporal_resolution,
patch_size=self.patch_size,
in_chans=self.in_chans,
embed_dim=self.embed_dim,
z_block_size=self.cfg.VIT.PATCH_SIZE_TEMP)
self.patch_embed_3d.proj.weight.data = torch.zeros_like(
self.patch_embed_3d.proj.weight.data)
# Number of patches
if self.video_input:
num_patches = self.patch_embed.num_patches * self.temporal_resolution
else:
num_patches = self.patch_embed.num_patches
self.num_patches = num_patches
# CLS token
self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
trunc_normal_(self.cls_token, std=.02)
# Positional embedding
self.pos_embed = nn.Parameter(
torch.zeros(1, self.patch_embed.num_patches + 1, self.embed_dim))
self.pos_drop = nn.Dropout(p=cfg.VIT.POS_DROPOUT)
trunc_normal_(self.pos_embed, std=.02)
if self.cfg.VIT.POS_EMBED == "joint":
self.st_embed = nn.Parameter(torch.zeros(1, num_patches + 1, self.embed_dim))
trunc_normal_(self.st_embed, std=.02)
elif self.cfg.VIT.POS_EMBED == "separate":
self.temp_embed = nn.Parameter(torch.zeros(1, self.temporal_resolution, self.embed_dim))
# Layer Blocks
dpr = [x.item() for x in torch.linspace(0, self.drop_path_rate, self.depth)]
if self.cfg.VIT.ATTN_LAYER == "divided":
self.blocks = nn.ModuleList([
vit_helper.DividedSpaceTimeBlock(
attn_type=cfg.VIT.ATTN_LAYER,
dim=self.embed_dim,
num_heads=self.num_heads,
mlp_ratio=self.mlp_ratio,
qkv_bias=self.qkv_bias,
drop=self.drop_rate,
attn_drop=self.attn_drop_rate,
drop_path=dpr[i],
norm_layer=norm_layer,
) for i in range(self.depth)
])
else:
self.blocks = nn.ModuleList([
vit_helper.Block(attn_type=cfg.VIT.ATTN_LAYER,
dim=self.embed_dim,
num_heads=self.num_heads,
mlp_ratio=self.mlp_ratio,
qkv_bias=self.qkv_bias,
drop=self.drop_rate,
attn_drop=self.attn_drop_rate,
drop_path=dpr[i],
norm_layer=norm_layer,
use_original_code=self.cfg.VIT.USE_ORIGINAL_TRAJ_ATTN_CODE)
for i in range(self.depth)
])
self.norm = norm_layer(self.embed_dim)
# MLP head
if self.use_mlp:
hidden_dim = self.embed_dim
if self.head_act == 'tanh':
# logging.info("Using TanH activation in MLP")
act = nn.Tanh()
elif self.head_act == 'gelu':
# logging.info("Using GELU activation in MLP")
act = nn.GELU()
else:
# logging.info("Using ReLU activation in MLP")
act = nn.ReLU()
self.pre_logits = nn.Sequential(
OrderedDict([
('fc', nn.Linear(self.embed_dim, hidden_dim)),
('act', act),
]))
else:
self.pre_logits = nn.Identity()
# Classifier Head
self.head_drop = nn.Dropout(p=self.head_dropout)
if isinstance(self.num_classes, (list, )) and len(self.num_classes) > 1:
for a, i in enumerate(range(len(self.num_classes))):
setattr(self, "head%d" % a, nn.Linear(self.embed_dim, self.num_classes[i]))
else:
self.head = nn.Linear(self.embed_dim,
self.num_classes) if self.num_classes > 0 else nn.Identity()
# Initialize weights
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
@torch.jit.ignore
def no_weight_decay(self):
if self.cfg.VIT.POS_EMBED == "joint":
return {'pos_embed', 'cls_token', 'st_embed'}
else:
return {'pos_embed', 'cls_token', 'temp_embed'}
def get_classifier(self):
return self.head
def reset_classifier(self, num_classes, global_pool=''):
self.num_classes = num_classes
self.head = (nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity())
def forward_features(self, x):
# if self.video_input:
# x = x[0]
B = x.shape[0]
# Tokenize input
# if self.cfg.VIT.PATCH_SIZE_TEMP > 1:
# for simplicity of mapping between content dimensions (input x) and token dims (after patching)
# we use the same trick as for AST (see modeling_ast.ASTModel.forward for the details):
# apply patching on input
x = self.patch_embed_3d(x)
tok_mask = None
# else:
# tok_mask = None
# # 2D tokenization
# if self.video_input:
# x = x.permute(0, 2, 1, 3, 4)
# (B, T, C, H, W) = x.shape
# x = x.reshape(B * T, C, H, W)
# x = self.patch_embed(x)
# if self.video_input:
# (B2, T2, D2) = x.shape
# x = x.reshape(B, T * T2, D2)
# Append CLS token
cls_tokens = self.cls_token.expand(B, -1, -1)
x = torch.cat((cls_tokens, x), dim=1)
# if tok_mask is not None:
# # prepend 1(=keep) to the mask to account for the CLS token as well
# tok_mask = torch.cat((torch.ones_like(tok_mask[:, [0]]), tok_mask), dim=1)
# Interpolate positinoal embeddings
# if self.cfg.DATA.TRAIN_CROP_SIZE != 224:
# pos_embed = self.pos_embed
# N = pos_embed.shape[1] - 1
# npatch = int((x.size(1) - 1) / self.temporal_resolution)
# class_emb = pos_embed[:, 0]
# pos_embed = pos_embed[:, 1:]
# dim = x.shape[-1]
# pos_embed = torch.nn.functional.interpolate(
# pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
# scale_factor=math.sqrt(npatch / N),
# mode='bicubic',
# )
# pos_embed = pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
# new_pos_embed = torch.cat((class_emb.unsqueeze(0), pos_embed), dim=1)
# else:
new_pos_embed = self.pos_embed
npatch = self.patch_embed.num_patches
# Add positional embeddings to input
if self.video_input:
if self.cfg.VIT.POS_EMBED == "separate":
cls_embed = self.pos_embed[:, 0, :].unsqueeze(1)
tile_pos_embed = new_pos_embed[:, 1:, :].repeat(1, self.temporal_resolution, 1)
tile_temporal_embed = self.temp_embed.repeat_interleave(npatch, 1)
total_pos_embed = tile_pos_embed + tile_temporal_embed
total_pos_embed = torch.cat([cls_embed, total_pos_embed], dim=1)
x = x + total_pos_embed
elif self.cfg.VIT.POS_EMBED == "joint":
x = x + self.st_embed
else:
# image input
x = x + new_pos_embed
# Apply positional dropout
x = self.pos_drop(x)
# Encoding using transformer layers
for i, blk in enumerate(self.blocks):
x = blk(x,
seq_len=npatch,
num_frames=self.temporal_resolution,
approx=self.cfg.VIT.APPROX_ATTN_TYPE,
num_landmarks=self.cfg.VIT.APPROX_ATTN_DIM,
tok_mask=tok_mask)
### v-iashin: I moved it to the forward pass
# x = self.norm(x)[:, 0]
# x = self.pre_logits(x)
###
return x, tok_mask
# def forward(self, x):
# x = self.forward_features(x)
# ### v-iashin: here. This should leave the same forward output as before
# x = self.norm(x)[:, 0]
# x = self.pre_logits(x)
# ###
# x = self.head_drop(x)
# if isinstance(self.num_classes, (list, )) and len(self.num_classes) > 1:
# output = []
# for head in range(len(self.num_classes)):
# x_out = getattr(self, "head%d" % head)(x)
# if not self.training:
# x_out = torch.nn.functional.softmax(x_out, dim=-1)
# output.append(x_out)
# return output
# else:
# x = self.head(x)
# if not self.training:
# x = torch.nn.functional.softmax(x, dim=-1)
# return x

399
postprocessing/mmaudio/ext/synchformer/vit_helper.py Normal file
View File

@@ -0,0 +1,399 @@
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
# Copyright 2020 Ross Wightman
# Modified Model definition
"""Video models."""
import math
import torch
import torch.nn as nn
from einops import rearrange, repeat
from timm.layers import to_2tuple
from torch import einsum
from torch.nn import functional as F
default_cfgs = {
'vit_1k':
'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_224-80ecf9dd.pth',
'vit_1k_large':
'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p16_224-4ee7a4dc.pth',
}
def qkv_attn(q, k, v, tok_mask: torch.Tensor = None):
sim = einsum('b i d, b j d -> b i j', q, k)
# apply masking if provided, tok_mask is (B*S*H, N): 1s - keep; sim is (B*S*H, H, N, N)
if tok_mask is not None:
BSH, N = tok_mask.shape
sim = sim.masked_fill(tok_mask.view(BSH, 1, N) == 0,
float('-inf')) # 1 - broadcasts across N
attn = sim.softmax(dim=-1)
out = einsum('b i j, b j d -> b i d', attn, v)
return out
class DividedAttention(nn.Module):
def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
super().__init__()
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = head_dim**-0.5
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.proj = nn.Linear(dim, dim)
# init to zeros
self.qkv.weight.data.fill_(0)
self.qkv.bias.data.fill_(0)
self.proj.weight.data.fill_(1)
self.proj.bias.data.fill_(0)
self.attn_drop = nn.Dropout(attn_drop)
self.proj_drop = nn.Dropout(proj_drop)
def forward(self, x, einops_from, einops_to, tok_mask: torch.Tensor = None, **einops_dims):
# num of heads variable
h = self.num_heads
# project x to q, k, v vaalues
q, k, v = self.qkv(x).chunk(3, dim=-1)
q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
if tok_mask is not None:
# replicate token mask across heads (b, n) -> (b, h, n) -> (b*h, n) -- same as qkv but w/o d
assert len(tok_mask.shape) == 2
tok_mask = tok_mask.unsqueeze(1).expand(-1, h, -1).reshape(-1, tok_mask.shape[1])
# Scale q
q *= self.scale
# Take out cls_q, cls_k, cls_v
(cls_q, q_), (cls_k, k_), (cls_v, v_) = map(lambda t: (t[:, 0:1], t[:, 1:]), (q, k, v))
# the same for masking
if tok_mask is not None:
cls_mask, mask_ = tok_mask[:, 0:1], tok_mask[:, 1:]
else:
cls_mask, mask_ = None, None
# let CLS token attend to key / values of all patches across time and space
cls_out = qkv_attn(cls_q, k, v, tok_mask=tok_mask)
# rearrange across time or space
q_, k_, v_ = map(lambda t: rearrange(t, f'{einops_from} -> {einops_to}', **einops_dims),
(q_, k_, v_))
# expand CLS token keys and values across time or space and concat
r = q_.shape[0] // cls_k.shape[0]
cls_k, cls_v = map(lambda t: repeat(t, 'b () d -> (b r) () d', r=r), (cls_k, cls_v))
k_ = torch.cat((cls_k, k_), dim=1)
v_ = torch.cat((cls_v, v_), dim=1)
# the same for masking (if provided)
if tok_mask is not None:
# since mask does not have the latent dim (d), we need to remove it from einops dims
mask_ = rearrange(mask_, f'{einops_from} -> {einops_to}'.replace(' d', ''),
**einops_dims)
cls_mask = repeat(cls_mask, 'b () -> (b r) ()',
r=r) # expand cls_mask across time or space
mask_ = torch.cat((cls_mask, mask_), dim=1)
# attention
out = qkv_attn(q_, k_, v_, tok_mask=mask_)
# merge back time or space
out = rearrange(out, f'{einops_to} -> {einops_from}', **einops_dims)
# concat back the cls token
out = torch.cat((cls_out, out), dim=1)
# merge back the heads
out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
## to out
x = self.proj(out)
x = self.proj_drop(x)
return x
class DividedSpaceTimeBlock(nn.Module):
def __init__(self,
dim=768,
num_heads=12,
attn_type='divided',
mlp_ratio=4.,
qkv_bias=False,
drop=0.,
attn_drop=0.,
drop_path=0.,
act_layer=nn.GELU,
norm_layer=nn.LayerNorm):
super().__init__()
self.einops_from_space = 'b (f n) d'
self.einops_to_space = '(b f) n d'
self.einops_from_time = 'b (f n) d'
self.einops_to_time = '(b n) f d'
self.norm1 = norm_layer(dim)
self.attn = DividedAttention(dim,
num_heads=num_heads,
qkv_bias=qkv_bias,
attn_drop=attn_drop,
proj_drop=drop)
self.timeattn = DividedAttention(dim,
num_heads=num_heads,
qkv_bias=qkv_bias,
attn_drop=attn_drop,
proj_drop=drop)
# self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
self.drop_path = nn.Identity()
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = Mlp(in_features=dim,
hidden_features=mlp_hidden_dim,
act_layer=act_layer,
drop=drop)
self.norm3 = norm_layer(dim)
def forward(self,
x,
seq_len=196,
num_frames=8,
approx='none',
num_landmarks=128,
tok_mask: torch.Tensor = None):
time_output = self.timeattn(self.norm3(x),
self.einops_from_time,
self.einops_to_time,
n=seq_len,
tok_mask=tok_mask)
time_residual = x + time_output
space_output = self.attn(self.norm1(time_residual),
self.einops_from_space,
self.einops_to_space,
f=num_frames,
tok_mask=tok_mask)
space_residual = time_residual + self.drop_path(space_output)
x = space_residual
x = x + self.drop_path(self.mlp(self.norm2(x)))
return x
class Mlp(nn.Module):
def __init__(self,
in_features,
hidden_features=None,
out_features=None,
act_layer=nn.GELU,
drop=0.):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
class PatchEmbed(nn.Module):
""" Image to Patch Embedding
"""
def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
super().__init__()
img_size = img_size if type(img_size) is tuple else to_2tuple(img_size)
patch_size = img_size if type(patch_size) is tuple else to_2tuple(patch_size)
num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
self.img_size = img_size
self.patch_size = patch_size
self.num_patches = num_patches
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
def forward(self, x):
B, C, H, W = x.shape
x = self.proj(x).flatten(2).transpose(1, 2)
return x
class PatchEmbed3D(nn.Module):
""" Image to Patch Embedding """
def __init__(self,
img_size=224,
temporal_resolution=4,
in_chans=3,
patch_size=16,
z_block_size=2,
embed_dim=768,
flatten=True):
super().__init__()
self.height = (img_size // patch_size)
self.width = (img_size // patch_size)
### v-iashin: these two are incorrect
# self.frames = (temporal_resolution // z_block_size)
# self.num_patches = self.height * self.width * self.frames
self.z_block_size = z_block_size
###
self.proj = nn.Conv3d(in_chans,
embed_dim,
kernel_size=(z_block_size, patch_size, patch_size),
stride=(z_block_size, patch_size, patch_size))
self.flatten = flatten
def forward(self, x):
B, C, T, H, W = x.shape
x = self.proj(x)
if self.flatten:
x = x.flatten(2).transpose(1, 2)
return x
class HeadMLP(nn.Module):
def __init__(self, n_input, n_classes, n_hidden=512, p=0.1):
super(HeadMLP, self).__init__()
self.n_input = n_input
self.n_classes = n_classes
self.n_hidden = n_hidden
if n_hidden is None:
# use linear classifier
self.block_forward = nn.Sequential(nn.Dropout(p=p),
nn.Linear(n_input, n_classes, bias=True))
else:
# use simple MLP classifier
self.block_forward = nn.Sequential(nn.Dropout(p=p),
nn.Linear(n_input, n_hidden, bias=True),
nn.BatchNorm1d(n_hidden), nn.ReLU(inplace=True),
nn.Dropout(p=p),
nn.Linear(n_hidden, n_classes, bias=True))
print(f"Dropout-NLP: {p}")
def forward(self, x):
return self.block_forward(x)
def _conv_filter(state_dict, patch_size=16):
""" convert patch embedding weight from manual patchify + linear proj to conv"""
out_dict = {}
for k, v in state_dict.items():
if 'patch_embed.proj.weight' in k:
v = v.reshape((v.shape[0], 3, patch_size, patch_size))
out_dict[k] = v
return out_dict
def adapt_input_conv(in_chans, conv_weight, agg='sum'):
conv_type = conv_weight.dtype
conv_weight = conv_weight.float()
O, I, J, K = conv_weight.shape
if in_chans == 1:
if I > 3:
assert conv_weight.shape[1] % 3 == 0
# For models with space2depth stems
conv_weight = conv_weight.reshape(O, I // 3, 3, J, K)
conv_weight = conv_weight.sum(dim=2, keepdim=False)
else:
if agg == 'sum':
print("Summing conv1 weights")
conv_weight = conv_weight.sum(dim=1, keepdim=True)
else:
print("Averaging conv1 weights")
conv_weight = conv_weight.mean(dim=1, keepdim=True)
elif in_chans != 3:
if I != 3:
raise NotImplementedError('Weight format not supported by conversion.')
else:
if agg == 'sum':
print("Summing conv1 weights")
repeat = int(math.ceil(in_chans / 3))
conv_weight = conv_weight.repeat(1, repeat, 1, 1)[:, :in_chans, :, :]
conv_weight *= (3 / float(in_chans))
else:
print("Averaging conv1 weights")
conv_weight = conv_weight.mean(dim=1, keepdim=True)
conv_weight = conv_weight.repeat(1, in_chans, 1, 1)
conv_weight = conv_weight.to(conv_type)
return conv_weight
def load_pretrained(model,
cfg=None,
num_classes=1000,
in_chans=3,
filter_fn=None,
strict=True,
progress=False):
# Load state dict
assert (f"{cfg.VIT.PRETRAINED_WEIGHTS} not in [vit_1k, vit_1k_large]")
state_dict = torch.hub.load_state_dict_from_url(url=default_cfgs[cfg.VIT.PRETRAINED_WEIGHTS])
if filter_fn is not None:
state_dict = filter_fn(state_dict)
input_convs = 'patch_embed.proj'
if input_convs is not None and in_chans != 3:
if isinstance(input_convs, str):
input_convs = (input_convs, )
for input_conv_name in input_convs:
weight_name = input_conv_name + '.weight'
try:
state_dict[weight_name] = adapt_input_conv(in_chans,
state_dict[weight_name],
agg='avg')
print(
f'Converted input conv {input_conv_name} pretrained weights from 3 to {in_chans} channel(s)'
)
except NotImplementedError as e:
del state_dict[weight_name]
strict = False
print(
f'Unable to convert pretrained {input_conv_name} weights, using random init for this layer.'
)
classifier_name = 'head'
label_offset = cfg.get('label_offset', 0)
pretrain_classes = 1000
if num_classes != pretrain_classes:
# completely discard fully connected if model num_classes doesn't match pretrained weights
del state_dict[classifier_name + '.weight']
del state_dict[classifier_name + '.bias']
strict = False
elif label_offset > 0:
# special case for pretrained weights with an extra background class in pretrained weights
classifier_weight = state_dict[classifier_name + '.weight']
state_dict[classifier_name + '.weight'] = classifier_weight[label_offset:]
classifier_bias = state_dict[classifier_name + '.bias']
state_dict[classifier_name + '.bias'] = classifier_bias[label_offset:]
loaded_state = state_dict
self_state = model.state_dict()
all_names = set(self_state.keys())
saved_names = set([])
for name, param in loaded_state.items():
param = param
if 'module.' in name:
name = name.replace('module.', '')
if name in self_state.keys() and param.shape == self_state[name].shape:
saved_names.add(name)
self_state[name].copy_(param)
else:
print(f"didnt load: {name} of shape: {param.shape}")
print("Missing Keys:")
print(all_names - saved_names)