codecsplat / src /model /encoder /unimatch /mv_transformer.py
multimodalart's picture
multimodalart HF Staff
CodecSplat demo
7405e85 verified
Raw History Blame Contribute Delete
22.2 kB
import torch
import torch.nn as nn
import torch.nn.functional as F
from .utils import split_feature, merge_splits
def gather_view_tensor(tensor, view_indices):
view_indices = view_indices.to(device=tensor.device, dtype=torch.long)
batch_indices = torch.arange(tensor.size(0), device=tensor.device)
batch_indices = batch_indices.reshape(
tensor.size(0), *([1] * (view_indices.dim() - 1))
)
return tensor[batch_indices, view_indices]
def non_self_view_indices(num_views, device):
view_ids = torch.arange(num_views, device=device)
return torch.stack(
[torch.cat((view_ids[:i], view_ids[i + 1 :])) for i in range(num_views)],
dim=0,
)
def target_view_indices(batch_size, num_views, nn_matrix, device):
if nn_matrix is not None:
return nn_matrix[..., 1:].to(device=device, dtype=torch.long)
return non_self_view_indices(num_views, device).unsqueeze(0).expand(
batch_size, -1, -1
)
def flatten_batch_view(tensor):
# Batch-major flatten: [B, V, ...] -> [B*V, ...].
return tensor.flatten(0, 1)
def flatten_view_batch(tensor):
# View-major flatten: [B, V, ...] -> [V*B, ...].
return tensor.transpose(0, 1).reshape(
tensor.size(0) * tensor.size(1), *tensor.shape[2:]
)
def unflatten_view_batch(tensor, batch_size, num_views):
# Inverse of flatten_view_batch.
return tensor.reshape(num_views, batch_size, *tensor.shape[1:]).transpose(0, 1)
def single_head_full_attention(q, k, v):
# q, k, v: [B, L, C]
assert q.dim() == k.dim() == v.dim() == 3
return F.scaled_dot_product_attention(
q.unsqueeze(1), k.unsqueeze(1), v.unsqueeze(1)
).squeeze(1)
def generate_shift_window_attn_mask(
input_resolution,
window_size_h,
window_size_w,
shift_size_h,
shift_size_w,
device=torch.device("cuda"),
):
# Ref: https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py
# Build the shifted-window attention mask used by Swin-style local attention.
h, w = input_resolution
img_mask = torch.zeros((1, h, w, 1)).to(device) # 1 H W 1
h_slices = (
slice(0, -window_size_h),
slice(-window_size_h, -shift_size_h),
slice(-shift_size_h, None),
)
w_slices = (
slice(0, -window_size_w),
slice(-window_size_w, -shift_size_w),
slice(-shift_size_w, None),
)
cnt = 0
for h in h_slices:
for w in w_slices:
img_mask[:, h, w, :] = cnt
cnt += 1
mask_windows = split_feature(
img_mask, num_splits=input_resolution[-1] // window_size_w, channel_last=True
)
mask_windows = mask_windows.view(-1, window_size_h * window_size_w)
attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(
attn_mask == 0, float(0.0)
)
return attn_mask
def single_head_split_window_attention(
q,
k,
v,
num_splits=1,
with_shift=False,
h=None,
w=None,
attn_mask=None,
):
# Ref: https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py
# q, k, v: [B, L, C] for self-attention or pairwise cross-attention.
# For multi-view cross-attention, q: [B, L, C], k/v: [B, M, L, C].
# Multi-view cross-attention, where M is the number of target views per query.
if not (q.dim() == k.dim() == v.dim() == 3):
assert k.dim() == v.dim() == 4
assert h is not None and w is not None
assert q.size(1) == h * w
m = k.size(1)
b, _, c = q.size()
b_new = b * num_splits * num_splits
window_size_h = h // num_splits
window_size_w = w // num_splits
q = q.view(b, h, w, c) # [B, H, W, C]
k = k.view(b, m, h, w, c) # [B, M, H, W, C]
v = v.view(b, m, h, w, c)
if with_shift:
assert attn_mask is not None # compute once
shift_size_h = window_size_h // 2
shift_size_w = window_size_w // 2
q = torch.roll(q, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2))
k = torch.roll(k, shifts=(-shift_size_h, -shift_size_w), dims=(2, 3))
v = torch.roll(v, shifts=(-shift_size_h, -shift_size_w), dims=(2, 3))
q = split_feature(
q, num_splits=num_splits, channel_last=True
) # [B*K*K, H/K, W/K, C]
k = split_feature(
k.permute(0, 2, 3, 4, 1).reshape(b, h, w, -1),
num_splits=num_splits,
channel_last=True,
) # [B*K*K, H/K, W/K, C*M]
v = split_feature(
v.permute(0, 2, 3, 4, 1).reshape(b, h, w, -1),
num_splits=num_splits,
channel_last=True,
) # [B*K*K, H/K, W/K, C*M]
k = (
k.view(b_new, h // num_splits, w // num_splits, c, m)
.permute(0, 1, 2, 4, 3)
.reshape(b_new, -1, c)
) # [B*K*K, H/K*W/K*M, C]
v = (
v.view(b_new, h // num_splits, w // num_splits, c, m)
.permute(0, 1, 2, 4, 3)
.reshape(b_new, -1, c)
) # [B*K*K, H/K*W/K*M, C]
attn_bias = None
if with_shift:
attn_bias = attn_mask.repeat(b, 1, m).unsqueeze(1).to(q.dtype)
out = F.scaled_dot_product_attention(
q.view(b_new, -1, c).unsqueeze(1),
k.unsqueeze(1),
v.unsqueeze(1),
attn_mask=attn_bias,
).squeeze(1)
out = merge_splits(
out.view(b_new, h // num_splits, w // num_splits, c),
num_splits=num_splits,
channel_last=True,
) # [B, H, W, C]
# shift back
if with_shift:
out = torch.roll(out, shifts=(shift_size_h, shift_size_w), dims=(1, 2))
out = out.view(b, -1, c)
else:
# Self-attention or pairwise cross-attention.
assert q.dim() == k.dim() == v.dim() == 3
assert h is not None and w is not None
assert q.size(1) == h * w
b, _, c = q.size()
b_new = b * num_splits * num_splits
window_size_h = h // num_splits
window_size_w = w // num_splits
q = q.view(b, h, w, c) # [B, H, W, C]
k = k.view(b, h, w, c)
v = v.view(b, h, w, c)
if with_shift:
assert attn_mask is not None # compute once
shift_size_h = window_size_h // 2
shift_size_w = window_size_w // 2
q = torch.roll(q, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2))
k = torch.roll(k, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2))
v = torch.roll(v, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2))
q = split_feature(
q, num_splits=num_splits, channel_last=True
) # [B*K*K, H/K, W/K, C]
k = split_feature(k, num_splits=num_splits, channel_last=True)
v = split_feature(v, num_splits=num_splits, channel_last=True)
attn_bias = None
if with_shift:
attn_bias = attn_mask.repeat(b, 1, 1).unsqueeze(1).to(q.dtype)
out = F.scaled_dot_product_attention(
q.view(b_new, -1, c).unsqueeze(1),
k.view(b_new, -1, c).unsqueeze(1),
v.view(b_new, -1, c).unsqueeze(1),
attn_mask=attn_bias,
).squeeze(1)
out = merge_splits(
out.view(b_new, h // num_splits, w // num_splits, c),
num_splits=num_splits,
channel_last=True,
) # [B, H, W, C]
# shift back
if with_shift:
out = torch.roll(out, shifts=(shift_size_h, shift_size_w), dims=(1, 2))
out = out.view(b, -1, c)
return out
def multi_head_split_window_attention(
q,
k,
v,
num_splits=1,
with_shift=False,
h=None,
w=None,
attn_mask=None,
num_head=1,
):
"""Windowed multi-head scaled dot-product attention.
Args:
q: [B, L, C]
k: [B, S, C]
v: [B, S, C]
Returns:
out: [B, L, C]
"""
assert h is not None and w is not None
assert q.size(1) == h * w
b, _, c = q.size()
b_new = b * num_splits * num_splits
window_size_h = h // num_splits
window_size_w = w // num_splits
q = q.view(b, h, w, c) # [B, H, W, C]
k = k.view(b, h, w, c)
v = v.view(b, h, w, c)
assert c % num_head == 0
if with_shift:
assert attn_mask is not None # compute once
shift_size_h = window_size_h // 2
shift_size_w = window_size_w // 2
q = torch.roll(q, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2))
k = torch.roll(k, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2))
v = torch.roll(v, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2))
q = split_feature(q, num_splits=num_splits) # [B*K*K, H/K, W/K, C]
k = split_feature(k, num_splits=num_splits)
v = split_feature(v, num_splits=num_splits)
# Split channels into attention heads for PyTorch SDPA.
head_dim = c // num_head
q = q.view(b_new, -1, num_head, head_dim).transpose(1, 2)
k = k.view(b_new, -1, num_head, head_dim).transpose(1, 2)
v = v.view(b_new, -1, num_head, head_dim).transpose(1, 2)
attn_bias = None
if with_shift:
attn_bias = attn_mask.unsqueeze(1).repeat(b, num_head, 1, 1).to(q.dtype)
out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_bias)
out = merge_splits(
out.transpose(1, 2).reshape(b_new, h // num_splits, w // num_splits, c),
num_splits=num_splits,
) # [B, H, W, C]
# shift back
if with_shift:
out = torch.roll(out, shifts=(shift_size_h, shift_size_w), dims=(1, 2))
out = out.view(b, -1, c)
return out
class TransformerLayer(nn.Module):
def __init__(
self,
d_model=256,
nhead=1,
attention_type="swin",
no_ffn=False,
ffn_dim_expansion=4,
with_shift=False,
add_per_view_attn=False,
**kwargs,
):
super(TransformerLayer, self).__init__()
self.dim = d_model
self.nhead = nhead
self.attention_type = attention_type
self.no_ffn = no_ffn
self.add_per_view_attn = add_per_view_attn
self.with_shift = with_shift
# Linear projections used by both single-head and multi-head attention.
self.q_proj = nn.Linear(d_model, d_model, bias=False)
self.k_proj = nn.Linear(d_model, d_model, bias=False)
self.v_proj = nn.Linear(d_model, d_model, bias=False)
self.merge = nn.Linear(d_model, d_model, bias=False)
self.norm1 = nn.LayerNorm(d_model)
# no ffn after self-attn, with ffn after cross-attn
if not self.no_ffn:
in_channels = d_model * 2
self.mlp = nn.Sequential(
nn.Linear(in_channels, in_channels * ffn_dim_expansion, bias=False),
nn.GELU(),
nn.Linear(in_channels * ffn_dim_expansion, d_model, bias=False),
)
self.norm2 = nn.LayerNorm(d_model)
def forward(
self,
source,
target,
height=None,
width=None,
shifted_window_attn_mask=None,
attn_num_splits=None,
**kwargs,
):
if "attn_type" in kwargs:
attn_type = kwargs["attn_type"]
else:
attn_type = self.attention_type
# source: [B, L, C]
# target: [B, L, C] for pairwise attention or [B, M, L, C] for multi-view.
query, key, value = source, target, target
# Project tokens before dispatching to local/global attention.
query = self.q_proj(query) # [B, L, C]
key = self.k_proj(key) # [B, L, C] or [B, M, L, C]
value = self.v_proj(value) # [B, L, C] or [B, M, L, C]
if attn_type == "swin" and attn_num_splits > 1:
if self.nhead > 1:
message = multi_head_split_window_attention(
query,
key,
value,
num_splits=attn_num_splits,
with_shift=self.with_shift,
h=height,
w=width,
attn_mask=shifted_window_attn_mask,
num_head=self.nhead,
)
else:
if self.add_per_view_attn:
assert query.dim() == 3 and key.dim() == 4 and value.dim() == 4
b, l, c = query.size()
query = query.unsqueeze(1).expand(
-1, key.size(1), -1, -1
).reshape(-1, l, c) # [B*M, L, C]
key = key.reshape(-1, l, c)
value = value.reshape(-1, l, c)
message = single_head_split_window_attention(
query,
key,
value,
num_splits=attn_num_splits,
with_shift=self.with_shift,
h=height,
w=width,
attn_mask=shifted_window_attn_mask,
)
# Sum per-target-view messages back to [B, L, C].
message = message.view(b, -1, l, c).sum(1)
else:
message = single_head_split_window_attention(
query,
key,
value,
num_splits=attn_num_splits,
with_shift=self.with_shift,
h=height,
w=width,
attn_mask=shifted_window_attn_mask,
)
else:
message = single_head_full_attention(query, key, value) # [B, L, C]
message = self.merge(message) # [B, L, C]
message = self.norm1(message)
if not self.no_ffn:
message = self.mlp(torch.cat([source, message], dim=-1))
message = self.norm2(message)
return source + message
class TransformerBlock(nn.Module):
"""self attention + cross attention + FFN"""
def __init__(
self,
d_model=256,
nhead=1,
attention_type="swin",
ffn_dim_expansion=4,
with_shift=False,
add_per_view_attn=False,
no_cross_attn=False,
**kwargs,
):
super(TransformerBlock, self).__init__()
self.no_cross_attn = no_cross_attn
if no_cross_attn:
self.self_attn = TransformerLayer(
d_model=d_model,
nhead=nhead,
attention_type=attention_type,
ffn_dim_expansion=ffn_dim_expansion,
with_shift=with_shift,
add_per_view_attn=add_per_view_attn,
)
else:
self.self_attn = TransformerLayer(
d_model=d_model,
nhead=nhead,
attention_type=attention_type,
no_ffn=True,
ffn_dim_expansion=ffn_dim_expansion,
with_shift=with_shift,
)
self.cross_attn_ffn = TransformerLayer(
d_model=d_model,
nhead=nhead,
attention_type=attention_type,
ffn_dim_expansion=ffn_dim_expansion,
with_shift=with_shift,
add_per_view_attn=add_per_view_attn,
)
def forward(
self,
source,
target,
height=None,
width=None,
shifted_window_attn_mask=None,
attn_num_splits=None,
**kwargs,
):
# source, target: [B, L, C]
# self attention
source = self.self_attn(
source,
source,
height=height,
width=width,
shifted_window_attn_mask=shifted_window_attn_mask,
attn_num_splits=attn_num_splits,
**kwargs,
)
if self.no_cross_attn:
return source
# cross attention and ffn
source = self.cross_attn_ffn(
source,
target,
height=height,
width=width,
shifted_window_attn_mask=shifted_window_attn_mask,
attn_num_splits=attn_num_splits,
**kwargs,
)
return source
def batch_features(features, nn_matrix=None):
# Construct transformer inputs from tensor views.
# features: [B, V, C, H, W] or [B, V, H*W, C].
# Returns view-major query/target tensors: [V*B, ...] and [V*B, M, ...].
assert features.dim() in (4, 5)
batch_size, num_views = features.shape[:2]
indices = target_view_indices(
batch_size, num_views, nn_matrix, features.device
)
q = flatten_view_batch(features)
kv = gather_view_tensor(features, indices)
kv = flatten_view_batch(kv)
return q, kv
class MultiViewFeatureTransformer(nn.Module):
def __init__(
self,
num_layers=6,
d_model=128,
nhead=1,
attention_type="swin",
ffn_dim_expansion=4,
add_per_view_attn=False,
no_cross_attn=False,
**kwargs,
):
super(MultiViewFeatureTransformer, self).__init__()
self.attention_type = attention_type
self.d_model = d_model
self.nhead = nhead
self.layers = nn.ModuleList(
[
TransformerBlock(
d_model=d_model,
nhead=nhead,
attention_type=attention_type,
ffn_dim_expansion=ffn_dim_expansion,
with_shift=(
True if attention_type == "swin" and i % 2 == 1 else False
),
add_per_view_attn=add_per_view_attn,
no_cross_attn=no_cross_attn,
)
for i in range(num_layers)
]
)
for p in self.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
# zero init layers beyond 6
if num_layers > 6:
for i in range(6, num_layers):
self.layers[i].self_attn.norm1.weight.data.zero_()
self.layers[i].self_attn.norm1.bias.data.zero_()
self.layers[i].cross_attn_ffn.norm2.weight.data.zero_()
self.layers[i].cross_attn_ffn.norm2.bias.data.zero_()
def forward(
self,
multi_view_features,
attn_num_splits=None,
**kwargs,
):
nn_matrix = kwargs.pop("nn_matrix", None)
# multi_view_features: [B, V, C, H, W].
assert multi_view_features.dim() == 5
b, num_views, c, h, w = multi_view_features.shape
assert self.d_model == c
if self.attention_type == "swin" and attn_num_splits > 1:
# global and refine use different number of splits
window_size_h = h // attn_num_splits
window_size_w = w // attn_num_splits
# compute attn mask once
shifted_window_attn_mask = generate_shift_window_attn_mask(
input_resolution=(h, w),
window_size_h=window_size_h,
window_size_w=window_size_w,
shift_size_h=window_size_h // 2,
shift_size_w=window_size_w // 2,
device=multi_view_features.device,
) # [K*K, H/K*W/K, H/K*W/K]
else:
shifted_window_attn_mask = None
# View-major query/target tensors: [V*B, C, H, W], [V*B, M, C, H, W].
concat0, concat1 = batch_features(multi_view_features, nn_matrix=nn_matrix)
concat0 = concat0.reshape(num_views * b, c, -1).permute(
0, 2, 1
) # [V*B, H*W, C]
c1_v = num_views - 1 if nn_matrix is None else nn_matrix.shape[-1] - 1
concat1 = concat1.reshape(num_views * b, c1_v, c, -1).permute(
0, 1, 3, 2
) # [V*B, M, H*W, C]
for i, layer in enumerate(self.layers):
concat0 = layer(
concat0,
concat1,
height=h,
width=w,
shifted_window_attn_mask=shifted_window_attn_mask,
attn_num_splits=attn_num_splits,
)
if i < len(self.layers) - 1:
# Rebuild target-view groups after each layer.
features = unflatten_view_batch(concat0, b, num_views)
concat0, concat1 = batch_features(features, nn_matrix=nn_matrix)
features = unflatten_view_batch(concat0, b, num_views)
features = features.reshape(b, num_views, h, w, c)
features = features.permute(0, 1, 4, 2, 3).contiguous()
return features
def batch_features_camera_parameters(
features,
intrinsics,
extrinsics,
nn_matrix=None,
):
# Construct batch-major inputs for plane-sweep warping.
# features: [B, V, C, H, W].
# intrinsics/extrinsics: [B, V, 3, 3] and [B, V, 4, 4].
# Returns [B*V, ...] query tensors and [B*V, M, ...] target tensors.
assert features.dim() == 5
assert intrinsics.shape[-2:] == (3, 3)
assert extrinsics.shape[-2:] == (4, 4)
batch_size, num_views, c, h, w = features.shape
indices = target_view_indices(
batch_size, num_views, nn_matrix, features.device
)
num_selected_views = indices.size(-1)
q = features
q_intrinsics = intrinsics
q_extrinsics = extrinsics
kv = gather_view_tensor(features, indices)
kv_intrinsics = gather_view_tensor(intrinsics, indices)
kv_extrinsics = gather_view_tensor(extrinsics, indices)
q = flatten_batch_view(q) # [BV, C, H, W]
q_intrinsics = flatten_batch_view(q_intrinsics) # [BV, 3, 3]
q_extrinsics = flatten_batch_view(q_extrinsics) # [BV, 4, 4]
kv = kv.reshape(-1, num_selected_views, c, h, w) # [BV, V-1, C, H, W]
kv_intrinsics = kv_intrinsics.reshape(-1, num_selected_views, 3, 3)
kv_extrinsics = kv_extrinsics.reshape(-1, num_selected_views, 4, 4)
return q, q_intrinsics, q_extrinsics, kv, kv_intrinsics, kv_extrinsics