release infer and demo
This commit is contained in:
pq-yang
@@ -0,0 +1,93 @@
|
||||
"""
|
||||
For computing auxiliary outputs for auxiliary losses
|
||||
"""
|
||||
from typing import Dict
|
||||
from omegaconf import DictConfig
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from matanyone2.model.group_modules import GConv2d
|
||||
from matanyone2.utils.tensor_utils import aggregate
|
||||
|
||||
|
||||
class LinearPredictor(nn.Module):
|
||||
def __init__(self, x_dim: int, pix_dim: int):
|
||||
super().__init__()
|
||||
self.projection = GConv2d(x_dim, pix_dim + 1, kernel_size=1)
|
||||
|
||||
def forward(self, pix_feat: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
|
||||
# pixel_feat: B*pix_dim*H*W
|
||||
# x: B*num_objects*x_dim*H*W
|
||||
num_objects = x.shape[1]
|
||||
x = self.projection(x)
|
||||
|
||||
pix_feat = pix_feat.unsqueeze(1).expand(-1, num_objects, -1, -1, -1)
|
||||
logits = (pix_feat * x[:, :, :-1]).sum(dim=2) + x[:, :, -1]
|
||||
return logits
|
||||
|
||||
|
||||
class DirectPredictor(nn.Module):
|
||||
def __init__(self, x_dim: int):
|
||||
super().__init__()
|
||||
self.projection = GConv2d(x_dim, 1, kernel_size=1)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
# x: B*num_objects*x_dim*H*W
|
||||
logits = self.projection(x).squeeze(2)
|
||||
return logits
|
||||
|
||||
|
||||
class AuxComputer(nn.Module):
|
||||
def __init__(self, cfg: DictConfig):
|
||||
super().__init__()
|
||||
|
||||
use_sensory_aux = cfg.model.aux_loss.sensory.enabled
|
||||
self.use_query_aux = cfg.model.aux_loss.query.enabled
|
||||
self.use_sensory_aux = use_sensory_aux
|
||||
|
||||
sensory_dim = cfg.model.sensory_dim
|
||||
embed_dim = cfg.model.embed_dim
|
||||
|
||||
if use_sensory_aux:
|
||||
self.sensory_aux = LinearPredictor(sensory_dim, embed_dim)
|
||||
|
||||
def _aggregate_with_selector(self, logits: torch.Tensor, selector: torch.Tensor) -> torch.Tensor:
|
||||
prob = torch.sigmoid(logits)
|
||||
if selector is not None:
|
||||
prob = prob * selector
|
||||
logits = aggregate(prob, dim=1)
|
||||
return logits
|
||||
|
||||
def forward(self, pix_feat: torch.Tensor, aux_input: Dict[str, torch.Tensor],
|
||||
selector: torch.Tensor, seg_pass=False) -> Dict[str, torch.Tensor]:
|
||||
sensory = aux_input['sensory']
|
||||
q_logits = aux_input['q_logits']
|
||||
|
||||
aux_output = {}
|
||||
aux_output['attn_mask'] = aux_input['attn_mask']
|
||||
|
||||
if self.use_sensory_aux:
|
||||
# B*num_objects*H*W
|
||||
logits = self.sensory_aux(pix_feat, sensory)
|
||||
aux_output['sensory_logits'] = self._aggregate_with_selector(logits, selector)
|
||||
if self.use_query_aux:
|
||||
# B*num_objects*num_levels*H*W
|
||||
aux_output['q_logits'] = self._aggregate_with_selector(
|
||||
torch.stack(q_logits, dim=2),
|
||||
selector.unsqueeze(2) if selector is not None else None)
|
||||
|
||||
return aux_output
|
||||
|
||||
def compute_mask(self, aux_input: Dict[str, torch.Tensor],
|
||||
selector: torch.Tensor) -> Dict[str, torch.Tensor]:
|
||||
# sensory = aux_input['sensory']
|
||||
q_logits = aux_input['q_logits']
|
||||
|
||||
aux_output = {}
|
||||
|
||||
# B*num_objects*num_levels*H*W
|
||||
aux_output['q_logits'] = self._aggregate_with_selector(
|
||||
torch.stack(q_logits, dim=2),
|
||||
selector.unsqueeze(2) if selector is not None else None)
|
||||
|
||||
return aux_output
|
||||
@@ -0,0 +1,366 @@
|
||||
"""
|
||||
big_modules.py - This file stores higher-level network blocks.
|
||||
|
||||
x - usually denotes features that are shared between objects.
|
||||
g - usually denotes features that are not shared between objects
|
||||
with an extra "num_objects" dimension (batch_size * num_objects * num_channels * H * W).
|
||||
|
||||
The trailing number of a variable usually denotes the stride
|
||||
"""
|
||||
|
||||
from typing import Iterable
|
||||
from omegaconf import DictConfig
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from matanyone2.model.group_modules import MainToGroupDistributor, GroupFeatureFusionBlock, GConv2d
|
||||
from matanyone2.model.utils import resnet
|
||||
from matanyone2.model.modules import SensoryDeepUpdater, SensoryUpdater_fullscale, DecoderFeatureProcessor, MaskUpsampleBlock
|
||||
from matanyone2.utils.device import safe_autocast
|
||||
|
||||
class UncertPred(nn.Module):
|
||||
def __init__(self, model_cfg: DictConfig):
|
||||
super().__init__()
|
||||
self.conv1x1_v2 = nn.Conv2d(model_cfg.pixel_dim*2 + 1 + model_cfg.value_dim, 64, kernel_size=1, stride=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(64)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.conv3x3 = nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1, groups=1, bias=False, dilation=1)
|
||||
self.bn2 = nn.BatchNorm2d(32)
|
||||
self.conv3x3_out = nn.Conv2d(32, 1, kernel_size=3, stride=1, padding=1, groups=1, bias=False, dilation=1)
|
||||
|
||||
def forward(self, last_frame_feat: torch.Tensor, cur_frame_feat: torch.Tensor, last_mask: torch.Tensor, mem_val_diff:torch.Tensor):
|
||||
last_mask = F.interpolate(last_mask, size=last_frame_feat.shape[-2:], mode='area')
|
||||
x = torch.cat([last_frame_feat, cur_frame_feat, last_mask, mem_val_diff], dim=1)
|
||||
x = self.conv1x1_v2(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
x = self.conv3x3(x)
|
||||
x = self.bn2(x)
|
||||
x = self.relu(x)
|
||||
x = self.conv3x3_out(x)
|
||||
return x
|
||||
|
||||
# override the default train() to freeze BN statistics
|
||||
def train(self, mode=True):
|
||||
self.training = False
|
||||
for module in self.children():
|
||||
module.train(False)
|
||||
return self
|
||||
|
||||
class PixelEncoder(nn.Module):
|
||||
def __init__(self, model_cfg: DictConfig):
|
||||
super().__init__()
|
||||
|
||||
self.is_resnet = 'resnet' in model_cfg.pixel_encoder.type
|
||||
# if model_cfg.pretrained_resnet is set in the model_cfg we get the value
|
||||
# else default to True
|
||||
is_pretrained_resnet = getattr(model_cfg,"pretrained_resnet",True)
|
||||
if self.is_resnet:
|
||||
if model_cfg.pixel_encoder.type == 'resnet18':
|
||||
network = resnet.resnet18(pretrained=is_pretrained_resnet)
|
||||
elif model_cfg.pixel_encoder.type == 'resnet50':
|
||||
network = resnet.resnet50(pretrained=is_pretrained_resnet)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
self.conv1 = network.conv1
|
||||
self.bn1 = network.bn1
|
||||
self.relu = network.relu
|
||||
self.maxpool = network.maxpool
|
||||
|
||||
self.res2 = network.layer1
|
||||
self.layer2 = network.layer2
|
||||
self.layer3 = network.layer3
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
def forward(self, x: torch.Tensor, seq_length=None) -> (torch.Tensor, torch.Tensor, torch.Tensor):
|
||||
f1 = x
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
f2 = x
|
||||
x = self.maxpool(x)
|
||||
f4 = self.res2(x)
|
||||
f8 = self.layer2(f4)
|
||||
f16 = self.layer3(f8)
|
||||
|
||||
return f16, f8, f4, f2, f1
|
||||
|
||||
# override the default train() to freeze BN statistics
|
||||
def train(self, mode=True):
|
||||
self.training = False
|
||||
for module in self.children():
|
||||
module.train(False)
|
||||
return self
|
||||
|
||||
|
||||
class KeyProjection(nn.Module):
|
||||
def __init__(self, model_cfg: DictConfig):
|
||||
super().__init__()
|
||||
in_dim = model_cfg.pixel_encoder.ms_dims[0]
|
||||
mid_dim = model_cfg.pixel_dim
|
||||
key_dim = model_cfg.key_dim
|
||||
|
||||
self.pix_feat_proj = nn.Conv2d(in_dim, mid_dim, kernel_size=1)
|
||||
self.key_proj = nn.Conv2d(mid_dim, key_dim, kernel_size=3, padding=1)
|
||||
# shrinkage
|
||||
self.d_proj = nn.Conv2d(mid_dim, 1, kernel_size=3, padding=1)
|
||||
# selection
|
||||
self.e_proj = nn.Conv2d(mid_dim, key_dim, kernel_size=3, padding=1)
|
||||
|
||||
nn.init.orthogonal_(self.key_proj.weight.data)
|
||||
nn.init.zeros_(self.key_proj.bias.data)
|
||||
|
||||
def forward(self, x: torch.Tensor, *, need_s: bool,
|
||||
need_e: bool) -> (torch.Tensor, torch.Tensor, torch.Tensor):
|
||||
x = self.pix_feat_proj(x)
|
||||
shrinkage = self.d_proj(x)**2 + 1 if (need_s) else None
|
||||
selection = torch.sigmoid(self.e_proj(x)) if (need_e) else None
|
||||
|
||||
return self.key_proj(x), shrinkage, selection
|
||||
|
||||
|
||||
class MaskEncoder(nn.Module):
|
||||
def __init__(self, model_cfg: DictConfig, single_object=False):
|
||||
super().__init__()
|
||||
pixel_dim = model_cfg.pixel_dim
|
||||
value_dim = model_cfg.value_dim
|
||||
sensory_dim = model_cfg.sensory_dim
|
||||
final_dim = model_cfg.mask_encoder.final_dim
|
||||
|
||||
self.single_object = single_object
|
||||
extra_dim = 1 if single_object else 2
|
||||
|
||||
# if model_cfg.pretrained_resnet is set in the model_cfg we get the value
|
||||
# else default to True
|
||||
is_pretrained_resnet = getattr(model_cfg,"pretrained_resnet",True)
|
||||
if model_cfg.mask_encoder.type == 'resnet18':
|
||||
network = resnet.resnet18(pretrained=is_pretrained_resnet, extra_dim=extra_dim)
|
||||
elif model_cfg.mask_encoder.type == 'resnet50':
|
||||
network = resnet.resnet50(pretrained=is_pretrained_resnet, extra_dim=extra_dim)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
self.conv1 = network.conv1
|
||||
self.bn1 = network.bn1
|
||||
self.relu = network.relu
|
||||
self.maxpool = network.maxpool
|
||||
|
||||
self.layer1 = network.layer1
|
||||
self.layer2 = network.layer2
|
||||
self.layer3 = network.layer3
|
||||
|
||||
self.distributor = MainToGroupDistributor()
|
||||
self.fuser = GroupFeatureFusionBlock(pixel_dim, final_dim, value_dim)
|
||||
|
||||
self.sensory_update = SensoryDeepUpdater(value_dim, sensory_dim)
|
||||
|
||||
def forward(self,
|
||||
image: torch.Tensor,
|
||||
pix_feat: torch.Tensor,
|
||||
sensory: torch.Tensor,
|
||||
masks: torch.Tensor,
|
||||
others: torch.Tensor,
|
||||
*,
|
||||
deep_update: bool = True,
|
||||
chunk_size: int = -1) -> (torch.Tensor, torch.Tensor):
|
||||
# ms_features are from the key encoder
|
||||
# we only use the first one (lowest resolution), following XMem
|
||||
if self.single_object:
|
||||
g = masks.unsqueeze(2)
|
||||
else:
|
||||
g = torch.stack([masks, others], dim=2)
|
||||
|
||||
g = self.distributor(image, g)
|
||||
|
||||
batch_size, num_objects = g.shape[:2]
|
||||
if chunk_size < 1 or chunk_size >= num_objects:
|
||||
chunk_size = num_objects
|
||||
fast_path = True
|
||||
new_sensory = sensory
|
||||
else:
|
||||
if deep_update:
|
||||
new_sensory = torch.empty_like(sensory)
|
||||
else:
|
||||
new_sensory = sensory
|
||||
fast_path = False
|
||||
|
||||
# chunk-by-chunk inference
|
||||
all_g = []
|
||||
for i in range(0, num_objects, chunk_size):
|
||||
if fast_path:
|
||||
g_chunk = g
|
||||
else:
|
||||
g_chunk = g[:, i:i + chunk_size]
|
||||
actual_chunk_size = g_chunk.shape[1]
|
||||
g_chunk = g_chunk.flatten(start_dim=0, end_dim=1)
|
||||
|
||||
g_chunk = self.conv1(g_chunk)
|
||||
g_chunk = self.bn1(g_chunk) # 1/2, 64
|
||||
g_chunk = self.maxpool(g_chunk) # 1/4, 64
|
||||
g_chunk = self.relu(g_chunk)
|
||||
|
||||
g_chunk = self.layer1(g_chunk) # 1/4
|
||||
g_chunk = self.layer2(g_chunk) # 1/8
|
||||
g_chunk = self.layer3(g_chunk) # 1/16
|
||||
|
||||
g_chunk = g_chunk.view(batch_size, actual_chunk_size, *g_chunk.shape[1:])
|
||||
g_chunk = self.fuser(pix_feat, g_chunk)
|
||||
all_g.append(g_chunk)
|
||||
if deep_update:
|
||||
if fast_path:
|
||||
new_sensory = self.sensory_update(g_chunk, sensory)
|
||||
else:
|
||||
new_sensory[:, i:i + chunk_size] = self.sensory_update(
|
||||
g_chunk, sensory[:, i:i + chunk_size])
|
||||
g = torch.cat(all_g, dim=1)
|
||||
|
||||
return g, new_sensory
|
||||
|
||||
# override the default train() to freeze BN statistics
|
||||
def train(self, mode=True):
|
||||
self.training = False
|
||||
for module in self.children():
|
||||
module.train(False)
|
||||
return self
|
||||
|
||||
|
||||
class PixelFeatureFuser(nn.Module):
|
||||
def __init__(self, model_cfg: DictConfig, single_object=False):
|
||||
super().__init__()
|
||||
value_dim = model_cfg.value_dim
|
||||
sensory_dim = model_cfg.sensory_dim
|
||||
pixel_dim = model_cfg.pixel_dim
|
||||
embed_dim = model_cfg.embed_dim
|
||||
self.single_object = single_object
|
||||
|
||||
self.fuser = GroupFeatureFusionBlock(pixel_dim, value_dim, embed_dim)
|
||||
if self.single_object:
|
||||
self.sensory_compress = GConv2d(sensory_dim + 1, value_dim, kernel_size=1)
|
||||
else:
|
||||
self.sensory_compress = GConv2d(sensory_dim + 2, value_dim, kernel_size=1)
|
||||
|
||||
def forward(self,
|
||||
pix_feat: torch.Tensor,
|
||||
pixel_memory: torch.Tensor,
|
||||
sensory_memory: torch.Tensor,
|
||||
last_mask: torch.Tensor,
|
||||
last_others: torch.Tensor,
|
||||
*,
|
||||
chunk_size: int = -1) -> torch.Tensor:
|
||||
batch_size, num_objects = pixel_memory.shape[:2]
|
||||
|
||||
if self.single_object:
|
||||
last_mask = last_mask.unsqueeze(2)
|
||||
else:
|
||||
last_mask = torch.stack([last_mask, last_others], dim=2)
|
||||
|
||||
if chunk_size < 1:
|
||||
chunk_size = num_objects
|
||||
|
||||
# chunk-by-chunk inference
|
||||
all_p16 = []
|
||||
for i in range(0, num_objects, chunk_size):
|
||||
sensory_readout = self.sensory_compress(
|
||||
torch.cat([sensory_memory[:, i:i + chunk_size], last_mask[:, i:i + chunk_size]], 2))
|
||||
p16 = pixel_memory[:, i:i + chunk_size] + sensory_readout
|
||||
p16 = self.fuser(pix_feat, p16)
|
||||
all_p16.append(p16)
|
||||
p16 = torch.cat(all_p16, dim=1)
|
||||
|
||||
return p16
|
||||
|
||||
|
||||
class MaskDecoder(nn.Module):
|
||||
def __init__(self, model_cfg: DictConfig):
|
||||
super().__init__()
|
||||
embed_dim = model_cfg.embed_dim
|
||||
sensory_dim = model_cfg.sensory_dim
|
||||
ms_image_dims = model_cfg.pixel_encoder.ms_dims
|
||||
up_dims = model_cfg.mask_decoder.up_dims
|
||||
|
||||
assert embed_dim == up_dims[0]
|
||||
|
||||
self.sensory_update = SensoryUpdater_fullscale([up_dims[0], up_dims[1], up_dims[2], up_dims[3], up_dims[4] + 1], sensory_dim,
|
||||
sensory_dim)
|
||||
|
||||
self.decoder_feat_proc = DecoderFeatureProcessor(ms_image_dims[1:], up_dims[:-1])
|
||||
self.up_16_8 = MaskUpsampleBlock(up_dims[0], up_dims[1])
|
||||
self.up_8_4 = MaskUpsampleBlock(up_dims[1], up_dims[2])
|
||||
# newly add for alpha matte
|
||||
self.up_4_2 = MaskUpsampleBlock(up_dims[2], up_dims[3])
|
||||
self.up_2_1 = MaskUpsampleBlock(up_dims[3], up_dims[4])
|
||||
|
||||
self.pred_seg = nn.Conv2d(up_dims[-1], 1, kernel_size=3, padding=1)
|
||||
self.pred_mat = nn.Conv2d(up_dims[-1], 1, kernel_size=3, padding=1)
|
||||
|
||||
def forward(self,
|
||||
ms_image_feat: Iterable[torch.Tensor],
|
||||
memory_readout: torch.Tensor,
|
||||
sensory: torch.Tensor,
|
||||
*,
|
||||
chunk_size: int = -1,
|
||||
update_sensory: bool = True,
|
||||
seg_pass: bool = False,
|
||||
last_mask=None,
|
||||
sigmoid_residual=False) -> (torch.Tensor, torch.Tensor):
|
||||
|
||||
batch_size, num_objects = memory_readout.shape[:2]
|
||||
f8, f4, f2, f1 = self.decoder_feat_proc(ms_image_feat[1:])
|
||||
if chunk_size < 1 or chunk_size >= num_objects:
|
||||
chunk_size = num_objects
|
||||
fast_path = True
|
||||
new_sensory = sensory
|
||||
else:
|
||||
if update_sensory:
|
||||
new_sensory = torch.empty_like(sensory)
|
||||
else:
|
||||
new_sensory = sensory
|
||||
fast_path = False
|
||||
|
||||
# chunk-by-chunk inference
|
||||
all_logits = []
|
||||
for i in range(0, num_objects, chunk_size):
|
||||
if fast_path:
|
||||
p16 = memory_readout
|
||||
else:
|
||||
p16 = memory_readout[:, i:i + chunk_size]
|
||||
actual_chunk_size = p16.shape[1]
|
||||
|
||||
p8 = self.up_16_8(p16, f8)
|
||||
p4 = self.up_8_4(p8, f4)
|
||||
p2 = self.up_4_2(p4, f2)
|
||||
p1 = self.up_2_1(p2, f1)
|
||||
with safe_autocast(enabled=False):
|
||||
if seg_pass:
|
||||
if last_mask is not None:
|
||||
res = self.pred_seg(F.relu(p1.flatten(start_dim=0, end_dim=1).float()))
|
||||
if sigmoid_residual:
|
||||
res = (torch.sigmoid(res) - 0.5) * 2 # regularization: (-1, 1) change on last mask
|
||||
logits = last_mask + res
|
||||
else:
|
||||
logits = self.pred_seg(F.relu(p1.flatten(start_dim=0, end_dim=1).float()))
|
||||
else:
|
||||
if last_mask is not None:
|
||||
res = self.pred_mat(F.relu(p1.flatten(start_dim=0, end_dim=1).float()))
|
||||
if sigmoid_residual:
|
||||
res = (torch.sigmoid(res) - 0.5) * 2 # regularization: (-1, 1) change on last mask
|
||||
logits = last_mask + res
|
||||
else:
|
||||
logits = self.pred_mat(F.relu(p1.flatten(start_dim=0, end_dim=1).float()))
|
||||
## SensoryUpdater_fullscale
|
||||
if update_sensory:
|
||||
p1 = torch.cat(
|
||||
[p1, logits.view(batch_size, actual_chunk_size, 1, *logits.shape[-2:])], 2)
|
||||
if fast_path:
|
||||
new_sensory = self.sensory_update([p16, p8, p4, p2, p1], sensory)
|
||||
else:
|
||||
new_sensory[:,
|
||||
i:i + chunk_size] = self.sensory_update([p16, p8, p4, p2, p1],
|
||||
sensory[:,
|
||||
i:i + chunk_size])
|
||||
all_logits.append(logits)
|
||||
logits = torch.cat(all_logits, dim=0)
|
||||
logits = logits.view(batch_size, num_objects, *logits.shape[-2:])
|
||||
|
||||
return new_sensory, logits
|
||||
@@ -0,0 +1,39 @@
|
||||
import math
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class CAResBlock(nn.Module):
|
||||
def __init__(self, in_dim: int, out_dim: int, residual: bool = True):
|
||||
super().__init__()
|
||||
self.residual = residual
|
||||
self.conv1 = nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1)
|
||||
self.conv2 = nn.Conv2d(out_dim, out_dim, kernel_size=3, padding=1)
|
||||
|
||||
t = int((abs(math.log2(out_dim)) + 1) // 2)
|
||||
k = t if t % 2 else t + 1
|
||||
self.pool = nn.AdaptiveAvgPool2d(1)
|
||||
self.conv = nn.Conv1d(1, 1, kernel_size=k, padding=(k - 1) // 2, bias=False)
|
||||
|
||||
if self.residual:
|
||||
if in_dim == out_dim:
|
||||
self.downsample = nn.Identity()
|
||||
else:
|
||||
self.downsample = nn.Conv2d(in_dim, out_dim, kernel_size=1)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
r = x
|
||||
x = self.conv1(F.relu(x))
|
||||
x = self.conv2(F.relu(x))
|
||||
|
||||
b, c = x.shape[:2]
|
||||
w = self.pool(x).view(b, 1, c)
|
||||
w = self.conv(w).transpose(-1, -2).unsqueeze(-1).sigmoid() # B*C*1*1
|
||||
|
||||
if self.residual:
|
||||
x = x * w + self.downsample(r)
|
||||
else:
|
||||
x = x * w
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,126 @@
|
||||
from typing import Optional
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from matanyone2.model.channel_attn import CAResBlock
|
||||
|
||||
def interpolate_groups(g: torch.Tensor, ratio: float, mode: str,
|
||||
align_corners: bool) -> torch.Tensor:
|
||||
batch_size, num_objects = g.shape[:2]
|
||||
g = F.interpolate(g.flatten(start_dim=0, end_dim=1),
|
||||
scale_factor=ratio,
|
||||
mode=mode,
|
||||
align_corners=align_corners)
|
||||
g = g.view(batch_size, num_objects, *g.shape[1:])
|
||||
return g
|
||||
|
||||
|
||||
def upsample_groups(g: torch.Tensor,
|
||||
ratio: float = 2,
|
||||
mode: str = 'bilinear',
|
||||
align_corners: bool = False) -> torch.Tensor:
|
||||
return interpolate_groups(g, ratio, mode, align_corners)
|
||||
|
||||
|
||||
def downsample_groups(g: torch.Tensor,
|
||||
ratio: float = 1 / 2,
|
||||
mode: str = 'area',
|
||||
align_corners: bool = None) -> torch.Tensor:
|
||||
return interpolate_groups(g, ratio, mode, align_corners)
|
||||
|
||||
|
||||
class GConv2d(nn.Conv2d):
|
||||
def forward(self, g: torch.Tensor) -> torch.Tensor:
|
||||
batch_size, num_objects = g.shape[:2]
|
||||
g = super().forward(g.flatten(start_dim=0, end_dim=1))
|
||||
return g.view(batch_size, num_objects, *g.shape[1:])
|
||||
|
||||
|
||||
class GroupResBlock(nn.Module):
|
||||
def __init__(self, in_dim: int, out_dim: int):
|
||||
super().__init__()
|
||||
|
||||
if in_dim == out_dim:
|
||||
self.downsample = nn.Identity()
|
||||
else:
|
||||
self.downsample = GConv2d(in_dim, out_dim, kernel_size=1)
|
||||
|
||||
self.conv1 = GConv2d(in_dim, out_dim, kernel_size=3, padding=1)
|
||||
self.conv2 = GConv2d(out_dim, out_dim, kernel_size=3, padding=1)
|
||||
|
||||
def forward(self, g: torch.Tensor) -> torch.Tensor:
|
||||
out_g = self.conv1(F.relu(g))
|
||||
out_g = self.conv2(F.relu(out_g))
|
||||
|
||||
g = self.downsample(g)
|
||||
|
||||
return out_g + g
|
||||
|
||||
|
||||
class MainToGroupDistributor(nn.Module):
|
||||
def __init__(self,
|
||||
x_transform: Optional[nn.Module] = None,
|
||||
g_transform: Optional[nn.Module] = None,
|
||||
method: str = 'cat',
|
||||
reverse_order: bool = False):
|
||||
super().__init__()
|
||||
|
||||
self.x_transform = x_transform
|
||||
self.g_transform = g_transform
|
||||
self.method = method
|
||||
self.reverse_order = reverse_order
|
||||
|
||||
def forward(self, x: torch.Tensor, g: torch.Tensor, skip_expand: bool = False) -> torch.Tensor:
|
||||
num_objects = g.shape[1]
|
||||
|
||||
if self.x_transform is not None:
|
||||
x = self.x_transform(x)
|
||||
|
||||
if self.g_transform is not None:
|
||||
g = self.g_transform(g)
|
||||
|
||||
if not skip_expand:
|
||||
x = x.unsqueeze(1).expand(-1, num_objects, -1, -1, -1)
|
||||
if self.method == 'cat':
|
||||
if self.reverse_order:
|
||||
g = torch.cat([g, x], 2)
|
||||
else:
|
||||
g = torch.cat([x, g], 2)
|
||||
elif self.method == 'add':
|
||||
g = x + g
|
||||
elif self.method == 'mulcat':
|
||||
g = torch.cat([x * g, g], dim=2)
|
||||
elif self.method == 'muladd':
|
||||
g = x * g + g
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
return g
|
||||
|
||||
|
||||
class GroupFeatureFusionBlock(nn.Module):
|
||||
def __init__(self, x_in_dim: int, g_in_dim: int, out_dim: int):
|
||||
super().__init__()
|
||||
|
||||
x_transform = nn.Conv2d(x_in_dim, out_dim, kernel_size=1)
|
||||
g_transform = GConv2d(g_in_dim, out_dim, kernel_size=1)
|
||||
|
||||
self.distributor = MainToGroupDistributor(x_transform=x_transform,
|
||||
g_transform=g_transform,
|
||||
method='add')
|
||||
self.block1 = CAResBlock(out_dim, out_dim)
|
||||
self.block2 = CAResBlock(out_dim, out_dim)
|
||||
|
||||
def forward(self, x: torch.Tensor, g: torch.Tensor) -> torch.Tensor:
|
||||
batch_size, num_objects = g.shape[:2]
|
||||
|
||||
g = self.distributor(x, g)
|
||||
|
||||
g = g.flatten(start_dim=0, end_dim=1)
|
||||
|
||||
g = self.block1(g)
|
||||
g = self.block2(g)
|
||||
|
||||
g = g.view(batch_size, num_objects, *g.shape[1:])
|
||||
|
||||
return g
|
||||
@@ -0,0 +1,338 @@
|
||||
from typing import List, Dict, Iterable, Tuple
|
||||
import logging
|
||||
from omegaconf import DictConfig
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from omegaconf import OmegaConf
|
||||
from huggingface_hub import PyTorchModelHubMixin
|
||||
|
||||
from matanyone2.model.big_modules import PixelEncoder, UncertPred, KeyProjection, MaskEncoder, PixelFeatureFuser, MaskDecoder
|
||||
from matanyone2.model.aux_modules import AuxComputer
|
||||
from matanyone2.model.utils.memory_utils import get_affinity, readout
|
||||
from matanyone2.model.transformer.object_transformer import QueryTransformer
|
||||
from matanyone2.model.transformer.object_summarizer import ObjectSummarizer
|
||||
from matanyone2.utils.tensor_utils import aggregate
|
||||
from matanyone2.utils.device import get_default_device, safe_autocast
|
||||
|
||||
device = get_default_device()
|
||||
|
||||
log = logging.getLogger()
|
||||
class MatAnyone2(nn.Module,
|
||||
PyTorchModelHubMixin,
|
||||
library_name="matanyone2",
|
||||
repo_url="https://github.com/pq-yang/MatAnyone2",
|
||||
coders={
|
||||
DictConfig: (
|
||||
lambda x: OmegaConf.to_container(x),
|
||||
lambda data: OmegaConf.create(data),
|
||||
)
|
||||
},
|
||||
):
|
||||
|
||||
def __init__(self, cfg: DictConfig, *, single_object=False):
|
||||
super().__init__()
|
||||
self.cfg = cfg
|
||||
model_cfg = cfg.model
|
||||
self.ms_dims = model_cfg.pixel_encoder.ms_dims
|
||||
self.key_dim = model_cfg.key_dim
|
||||
self.value_dim = model_cfg.value_dim
|
||||
self.sensory_dim = model_cfg.sensory_dim
|
||||
self.pixel_dim = model_cfg.pixel_dim
|
||||
self.embed_dim = model_cfg.embed_dim
|
||||
self.single_object = single_object
|
||||
|
||||
log.info(f'Single object: {self.single_object}')
|
||||
|
||||
self.pixel_encoder = PixelEncoder(model_cfg)
|
||||
self.pix_feat_proj = nn.Conv2d(self.ms_dims[0], self.pixel_dim, kernel_size=1)
|
||||
self.key_proj = KeyProjection(model_cfg)
|
||||
self.mask_encoder = MaskEncoder(model_cfg, single_object=single_object)
|
||||
self.mask_decoder = MaskDecoder(model_cfg)
|
||||
self.pixel_fuser = PixelFeatureFuser(model_cfg, single_object=single_object)
|
||||
self.object_transformer = QueryTransformer(model_cfg)
|
||||
self.object_summarizer = ObjectSummarizer(model_cfg)
|
||||
self.aux_computer = AuxComputer(cfg)
|
||||
self.temp_sparity = UncertPred(model_cfg)
|
||||
|
||||
self.register_buffer("pixel_mean", torch.Tensor(model_cfg.pixel_mean).view(-1, 1, 1), False)
|
||||
self.register_buffer("pixel_std", torch.Tensor(model_cfg.pixel_std).view(-1, 1, 1), False)
|
||||
|
||||
def _get_others(self, masks: torch.Tensor) -> torch.Tensor:
|
||||
# for each object, return the sum of masks of all other objects
|
||||
if self.single_object:
|
||||
return None
|
||||
|
||||
num_objects = masks.shape[1]
|
||||
if num_objects >= 1:
|
||||
others = (masks.sum(dim=1, keepdim=True) - masks).clamp(0, 1)
|
||||
else:
|
||||
others = torch.zeros_like(masks)
|
||||
return others
|
||||
|
||||
def pred_uncertainty(self, last_pix_feat: torch.Tensor, cur_pix_feat: torch.Tensor, last_mask: torch.Tensor, mem_val_diff:torch.Tensor):
|
||||
logits = self.temp_sparity(last_frame_feat=last_pix_feat,
|
||||
cur_frame_feat=cur_pix_feat,
|
||||
last_mask=last_mask,
|
||||
mem_val_diff=mem_val_diff)
|
||||
|
||||
prob = torch.sigmoid(logits)
|
||||
mask = (prob > 0) + 0
|
||||
|
||||
uncert_output = {"logits": logits,
|
||||
"prob": prob,
|
||||
"mask": mask}
|
||||
|
||||
return uncert_output
|
||||
|
||||
def encode_image(self, image: torch.Tensor, seq_length=None, last_feats=None) -> (Iterable[torch.Tensor], torch.Tensor): # type: ignore
|
||||
self.pixel_mean = self.pixel_mean.to(device)
|
||||
self.pixel_std = self.pixel_std.to(device)
|
||||
image = (image - self.pixel_mean) / self.pixel_std
|
||||
ms_image_feat = self.pixel_encoder(image, seq_length) # f16, f8, f4, f2, f1
|
||||
return ms_image_feat, self.pix_feat_proj(ms_image_feat[0])
|
||||
|
||||
def encode_mask(
|
||||
self,
|
||||
image: torch.Tensor,
|
||||
ms_features: List[torch.Tensor],
|
||||
sensory: torch.Tensor,
|
||||
masks: torch.Tensor,
|
||||
*,
|
||||
deep_update: bool = True,
|
||||
chunk_size: int = -1,
|
||||
need_weights: bool = False) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
image = (image - self.pixel_mean) / self.pixel_std
|
||||
others = self._get_others(masks)
|
||||
mask_value, new_sensory = self.mask_encoder(image,
|
||||
ms_features,
|
||||
sensory,
|
||||
masks,
|
||||
others,
|
||||
deep_update=deep_update,
|
||||
chunk_size=chunk_size)
|
||||
object_summaries, object_logits = self.object_summarizer(masks, mask_value, need_weights)
|
||||
return mask_value, new_sensory, object_summaries, object_logits
|
||||
|
||||
def transform_key(self,
|
||||
final_pix_feat: torch.Tensor,
|
||||
*,
|
||||
need_sk: bool = True,
|
||||
need_ek: bool = True) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
key, shrinkage, selection = self.key_proj(final_pix_feat, need_s=need_sk, need_e=need_ek)
|
||||
return key, shrinkage, selection
|
||||
|
||||
# Used in training only.
|
||||
# This step is replaced by MemoryManager in test time
|
||||
def read_memory(self, query_key: torch.Tensor, query_selection: torch.Tensor,
|
||||
memory_key: torch.Tensor, memory_shrinkage: torch.Tensor,
|
||||
msk_value: torch.Tensor, obj_memory: torch.Tensor, pix_feat: torch.Tensor,
|
||||
sensory: torch.Tensor, last_mask: torch.Tensor,
|
||||
selector: torch.Tensor, uncert_output=None, seg_pass=False,
|
||||
last_pix_feat=None, last_pred_mask=None) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
|
||||
"""
|
||||
query_key : B * CK * H * W
|
||||
query_selection : B * CK * H * W
|
||||
memory_key : B * CK * T * H * W
|
||||
memory_shrinkage: B * 1 * T * H * W
|
||||
msk_value : B * num_objects * CV * T * H * W
|
||||
obj_memory : B * num_objects * T * num_summaries * C
|
||||
pixel_feature : B * C * H * W
|
||||
"""
|
||||
batch_size, num_objects = msk_value.shape[:2]
|
||||
|
||||
uncert_mask = uncert_output["mask"] if uncert_output is not None else None
|
||||
|
||||
# read using visual attention
|
||||
with safe_autocast(enabled=False):
|
||||
affinity = get_affinity(memory_key.float(), memory_shrinkage.float(), query_key.float(),
|
||||
query_selection.float(), uncert_mask=uncert_mask)
|
||||
|
||||
msk_value = msk_value.flatten(start_dim=1, end_dim=2).float()
|
||||
|
||||
# B * (num_objects*CV) * H * W
|
||||
pixel_readout = readout(affinity, msk_value, uncert_mask)
|
||||
pixel_readout = pixel_readout.view(batch_size, num_objects, self.value_dim,
|
||||
*pixel_readout.shape[-2:])
|
||||
|
||||
uncert_output = self.pred_uncertainty(last_pix_feat, pix_feat, last_pred_mask, pixel_readout[:,0]-msk_value[:,:,-1])
|
||||
uncert_prob = uncert_output["prob"].unsqueeze(1) # b n 1 h w
|
||||
pixel_readout = pixel_readout*uncert_prob + msk_value[:,:,-1].unsqueeze(1)*(1-uncert_prob)
|
||||
|
||||
pixel_readout = self.pixel_fusion(pix_feat, pixel_readout, sensory, last_mask)
|
||||
|
||||
|
||||
# read from query transformer
|
||||
mem_readout, aux_features = self.readout_query(pixel_readout, obj_memory, selector=selector, seg_pass=seg_pass)
|
||||
|
||||
aux_output = {
|
||||
'sensory': sensory,
|
||||
'q_logits': aux_features['logits'] if aux_features else None,
|
||||
'attn_mask': aux_features['attn_mask'] if aux_features else None,
|
||||
}
|
||||
|
||||
return mem_readout, aux_output, uncert_output
|
||||
|
||||
def read_first_frame_memory(self, pixel_readout,
|
||||
obj_memory: torch.Tensor, pix_feat: torch.Tensor,
|
||||
sensory: torch.Tensor, last_mask: torch.Tensor,
|
||||
selector: torch.Tensor, seg_pass=False) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
|
||||
"""
|
||||
query_key : B * CK * H * W
|
||||
query_selection : B * CK * H * W
|
||||
memory_key : B * CK * T * H * W
|
||||
memory_shrinkage: B * 1 * T * H * W
|
||||
msk_value : B * num_objects * CV * T * H * W
|
||||
obj_memory : B * num_objects * T * num_summaries * C
|
||||
pixel_feature : B * C * H * W
|
||||
"""
|
||||
|
||||
pixel_readout = self.pixel_fusion(pix_feat, pixel_readout, sensory, last_mask)
|
||||
|
||||
# read from query transformer
|
||||
mem_readout, aux_features = self.readout_query(pixel_readout, obj_memory, selector=selector, seg_pass=seg_pass)
|
||||
|
||||
aux_output = {
|
||||
'sensory': sensory,
|
||||
'q_logits': aux_features['logits'] if aux_features else None,
|
||||
'attn_mask': aux_features['attn_mask'] if aux_features else None,
|
||||
}
|
||||
|
||||
return mem_readout, aux_output
|
||||
|
||||
def pixel_fusion(self,
|
||||
pix_feat: torch.Tensor,
|
||||
pixel: torch.Tensor,
|
||||
sensory: torch.Tensor,
|
||||
last_mask: torch.Tensor,
|
||||
*,
|
||||
chunk_size: int = -1) -> torch.Tensor:
|
||||
last_mask = F.interpolate(last_mask, size=sensory.shape[-2:], mode='area')
|
||||
last_others = self._get_others(last_mask)
|
||||
fused = self.pixel_fuser(pix_feat,
|
||||
pixel,
|
||||
sensory,
|
||||
last_mask,
|
||||
last_others,
|
||||
chunk_size=chunk_size)
|
||||
return fused
|
||||
|
||||
def readout_query(self,
|
||||
pixel_readout,
|
||||
obj_memory,
|
||||
*,
|
||||
selector=None,
|
||||
need_weights=False,
|
||||
seg_pass=False) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
|
||||
return self.object_transformer(pixel_readout,
|
||||
obj_memory,
|
||||
selector=selector,
|
||||
need_weights=need_weights,
|
||||
seg_pass=seg_pass)
|
||||
|
||||
def segment(self,
|
||||
ms_image_feat: List[torch.Tensor],
|
||||
memory_readout: torch.Tensor,
|
||||
sensory: torch.Tensor,
|
||||
*,
|
||||
selector: bool = None,
|
||||
chunk_size: int = -1,
|
||||
update_sensory: bool = True,
|
||||
seg_pass: bool = False,
|
||||
clamp_mat: bool = True,
|
||||
last_mask=None,
|
||||
sigmoid_residual=False,
|
||||
seg_mat=False) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
multi_scale_features is from the key encoder for skip-connection
|
||||
memory_readout is from working/long-term memory
|
||||
sensory is the sensory memory
|
||||
last_mask is the mask from the last frame, supplementing sensory memory
|
||||
selector is 1 if an object exists, and 0 otherwise. We use it to filter padded objects
|
||||
during training.
|
||||
"""
|
||||
#### use mat head for seg data
|
||||
if seg_mat:
|
||||
assert seg_pass
|
||||
seg_pass = False
|
||||
####
|
||||
sensory, logits = self.mask_decoder(ms_image_feat,
|
||||
memory_readout,
|
||||
sensory,
|
||||
chunk_size=chunk_size,
|
||||
update_sensory=update_sensory,
|
||||
seg_pass = seg_pass,
|
||||
last_mask=last_mask,
|
||||
sigmoid_residual=sigmoid_residual)
|
||||
if seg_pass:
|
||||
prob = torch.sigmoid(logits)
|
||||
if selector is not None:
|
||||
prob = prob * selector
|
||||
|
||||
# Softmax over all objects[]
|
||||
logits = aggregate(prob, dim=1)
|
||||
prob = F.softmax(logits, dim=1)
|
||||
else:
|
||||
if clamp_mat:
|
||||
logits = logits.clamp(0.0, 1.0)
|
||||
logits = torch.cat([torch.prod(1 - logits, dim=1, keepdim=True), logits], 1)
|
||||
prob = logits
|
||||
|
||||
return sensory, logits, prob
|
||||
|
||||
def compute_aux(self, pix_feat: torch.Tensor, aux_inputs: Dict[str, torch.Tensor],
|
||||
selector: torch.Tensor, seg_pass=False) -> Dict[str, torch.Tensor]:
|
||||
return self.aux_computer(pix_feat, aux_inputs, selector, seg_pass=seg_pass)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
raise NotImplementedError
|
||||
|
||||
def load_weights(self, src_dict, init_as_zero_if_needed=False) -> None:
|
||||
if not self.single_object:
|
||||
# Map single-object weight to multi-object weight (4->5 out channels in conv1)
|
||||
for k in list(src_dict.keys()):
|
||||
if k == 'mask_encoder.conv1.weight':
|
||||
if src_dict[k].shape[1] == 4:
|
||||
log.info(f'Converting {k} from single object to multiple objects.')
|
||||
pads = torch.zeros((64, 1, 7, 7), device=src_dict[k].device)
|
||||
if not init_as_zero_if_needed:
|
||||
nn.init.orthogonal_(pads)
|
||||
log.info(f'Randomly initialized padding for {k}.')
|
||||
else:
|
||||
log.info(f'Zero-initialized padding for {k}.')
|
||||
src_dict[k] = torch.cat([src_dict[k], pads], 1)
|
||||
elif k == 'pixel_fuser.sensory_compress.weight':
|
||||
if src_dict[k].shape[1] == self.sensory_dim + 1:
|
||||
log.info(f'Converting {k} from single object to multiple objects.')
|
||||
pads = torch.zeros((self.value_dim, 1, 1, 1), device=src_dict[k].device)
|
||||
if not init_as_zero_if_needed:
|
||||
nn.init.orthogonal_(pads)
|
||||
log.info(f'Randomly initialized padding for {k}.')
|
||||
else:
|
||||
log.info(f'Zero-initialized padding for {k}.')
|
||||
src_dict[k] = torch.cat([src_dict[k], pads], 1)
|
||||
elif self.single_object:
|
||||
"""
|
||||
If the model is multiple-object and we are training in single-object,
|
||||
we strip the last channel of conv1.
|
||||
This is not supposed to happen in standard training except when users are trying to
|
||||
finetune a trained model with single object datasets.
|
||||
"""
|
||||
if src_dict['mask_encoder.conv1.weight'].shape[1] == 5:
|
||||
log.warning('Converting mask_encoder.conv1.weight from multiple objects to single object.'
|
||||
'This is not supposed to happen in standard training.')
|
||||
src_dict['mask_encoder.conv1.weight'] = src_dict['mask_encoder.conv1.weight'][:, :-1]
|
||||
src_dict['pixel_fuser.sensory_compress.weight'] = src_dict['pixel_fuser.sensory_compress.weight'][:, :-1]
|
||||
|
||||
for k in src_dict:
|
||||
if k not in self.state_dict():
|
||||
log.info(f'Key {k} found in src_dict but not in self.state_dict()!!!')
|
||||
for k in self.state_dict():
|
||||
if k not in src_dict:
|
||||
log.info(f'Key {k} found in self.state_dict() but not in src_dict!!!')
|
||||
|
||||
self.load_state_dict(src_dict, strict=False)
|
||||
|
||||
@property
|
||||
def device(self) -> torch.device:
|
||||
return self.pixel_mean.device
|
||||
@@ -0,0 +1,150 @@
|
||||
from typing import List, Iterable
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from matanyone2.model.group_modules import MainToGroupDistributor, GroupResBlock, upsample_groups, GConv2d, downsample_groups
|
||||
from matanyone2.utils.device import safe_autocast
|
||||
|
||||
|
||||
class UpsampleBlock(nn.Module):
|
||||
def __init__(self, in_dim: int, out_dim: int, scale_factor: int = 2):
|
||||
super().__init__()
|
||||
self.out_conv = ResBlock(in_dim, out_dim)
|
||||
self.scale_factor = scale_factor
|
||||
|
||||
def forward(self, in_g: torch.Tensor, skip_f: torch.Tensor) -> torch.Tensor:
|
||||
g = F.interpolate(in_g,
|
||||
scale_factor=self.scale_factor,
|
||||
mode='bilinear')
|
||||
g = self.out_conv(g)
|
||||
g = g + skip_f
|
||||
return g
|
||||
|
||||
class MaskUpsampleBlock(nn.Module):
|
||||
def __init__(self, in_dim: int, out_dim: int, scale_factor: int = 2):
|
||||
super().__init__()
|
||||
self.distributor = MainToGroupDistributor(method='add')
|
||||
self.out_conv = GroupResBlock(in_dim, out_dim)
|
||||
self.scale_factor = scale_factor
|
||||
|
||||
def forward(self, in_g: torch.Tensor, skip_f: torch.Tensor) -> torch.Tensor:
|
||||
g = upsample_groups(in_g, ratio=self.scale_factor)
|
||||
g = self.distributor(skip_f, g)
|
||||
g = self.out_conv(g)
|
||||
return g
|
||||
|
||||
|
||||
class DecoderFeatureProcessor(nn.Module):
|
||||
def __init__(self, decoder_dims: List[int], out_dims: List[int]):
|
||||
super().__init__()
|
||||
self.transforms = nn.ModuleList([
|
||||
nn.Conv2d(d_dim, p_dim, kernel_size=1) for d_dim, p_dim in zip(decoder_dims, out_dims)
|
||||
])
|
||||
|
||||
def forward(self, multi_scale_features: Iterable[torch.Tensor]) -> List[torch.Tensor]:
|
||||
outputs = [func(x) for x, func in zip(multi_scale_features, self.transforms)]
|
||||
return outputs
|
||||
|
||||
|
||||
# @torch.jit.script
|
||||
def _recurrent_update(h: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
|
||||
# h: batch_size * num_objects * hidden_dim * h * w
|
||||
# values: batch_size * num_objects * (hidden_dim*3) * h * w
|
||||
dim = values.shape[2] // 3
|
||||
forget_gate = torch.sigmoid(values[:, :, :dim])
|
||||
update_gate = torch.sigmoid(values[:, :, dim:dim * 2])
|
||||
new_value = torch.tanh(values[:, :, dim * 2:])
|
||||
new_h = forget_gate * h * (1 - update_gate) + update_gate * new_value
|
||||
return new_h
|
||||
|
||||
|
||||
class SensoryUpdater_fullscale(nn.Module):
|
||||
# Used in the decoder, multi-scale feature + GRU
|
||||
def __init__(self, g_dims: List[int], mid_dim: int, sensory_dim: int):
|
||||
super().__init__()
|
||||
self.g16_conv = GConv2d(g_dims[0], mid_dim, kernel_size=1)
|
||||
self.g8_conv = GConv2d(g_dims[1], mid_dim, kernel_size=1)
|
||||
self.g4_conv = GConv2d(g_dims[2], mid_dim, kernel_size=1)
|
||||
self.g2_conv = GConv2d(g_dims[3], mid_dim, kernel_size=1)
|
||||
self.g1_conv = GConv2d(g_dims[4], mid_dim, kernel_size=1)
|
||||
|
||||
self.transform = GConv2d(mid_dim + sensory_dim, sensory_dim * 3, kernel_size=3, padding=1)
|
||||
|
||||
nn.init.xavier_normal_(self.transform.weight)
|
||||
|
||||
def forward(self, g: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
|
||||
g = self.g16_conv(g[0]) + self.g8_conv(downsample_groups(g[1], ratio=1/2)) + \
|
||||
self.g4_conv(downsample_groups(g[2], ratio=1/4)) + \
|
||||
self.g2_conv(downsample_groups(g[3], ratio=1/8)) + \
|
||||
self.g1_conv(downsample_groups(g[4], ratio=1/16))
|
||||
|
||||
with safe_autocast(enabled=False):
|
||||
g = g.float()
|
||||
h = h.float()
|
||||
values = self.transform(torch.cat([g, h], dim=2))
|
||||
new_h = _recurrent_update(h, values)
|
||||
|
||||
return new_h
|
||||
|
||||
class SensoryUpdater(nn.Module):
|
||||
# Used in the decoder, multi-scale feature + GRU
|
||||
def __init__(self, g_dims: List[int], mid_dim: int, sensory_dim: int):
|
||||
super().__init__()
|
||||
self.g16_conv = GConv2d(g_dims[0], mid_dim, kernel_size=1)
|
||||
self.g8_conv = GConv2d(g_dims[1], mid_dim, kernel_size=1)
|
||||
self.g4_conv = GConv2d(g_dims[2], mid_dim, kernel_size=1)
|
||||
|
||||
self.transform = GConv2d(mid_dim + sensory_dim, sensory_dim * 3, kernel_size=3, padding=1)
|
||||
|
||||
nn.init.xavier_normal_(self.transform.weight)
|
||||
|
||||
def forward(self, g: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
|
||||
g = self.g16_conv(g[0]) + self.g8_conv(downsample_groups(g[1], ratio=1/2)) + \
|
||||
self.g4_conv(downsample_groups(g[2], ratio=1/4))
|
||||
|
||||
with safe_autocast(enabled=False):
|
||||
g = g.float()
|
||||
h = h.float()
|
||||
values = self.transform(torch.cat([g, h], dim=2))
|
||||
new_h = _recurrent_update(h, values)
|
||||
|
||||
return new_h
|
||||
|
||||
|
||||
class SensoryDeepUpdater(nn.Module):
|
||||
def __init__(self, f_dim: int, sensory_dim: int):
|
||||
super().__init__()
|
||||
self.transform = GConv2d(f_dim + sensory_dim, sensory_dim * 3, kernel_size=3, padding=1)
|
||||
|
||||
nn.init.xavier_normal_(self.transform.weight)
|
||||
|
||||
def forward(self, g: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
|
||||
with safe_autocast(enabled=False):
|
||||
g = g.float()
|
||||
h = h.float()
|
||||
values = self.transform(torch.cat([g, h], dim=2))
|
||||
new_h = _recurrent_update(h, values)
|
||||
|
||||
return new_h
|
||||
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
def __init__(self, in_dim: int, out_dim: int):
|
||||
super().__init__()
|
||||
|
||||
if in_dim == out_dim:
|
||||
self.downsample = nn.Identity()
|
||||
else:
|
||||
self.downsample = nn.Conv2d(in_dim, out_dim, kernel_size=1)
|
||||
|
||||
self.conv1 = nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1)
|
||||
self.conv2 = nn.Conv2d(out_dim, out_dim, kernel_size=3, padding=1)
|
||||
|
||||
def forward(self, g: torch.Tensor) -> torch.Tensor:
|
||||
out_g = self.conv1(F.relu(g))
|
||||
out_g = self.conv2(F.relu(out_g))
|
||||
|
||||
g = self.downsample(g)
|
||||
|
||||
return out_g + g
|
||||
@@ -0,0 +1,90 @@
|
||||
from typing import Optional
|
||||
from omegaconf import DictConfig
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from matanyone2.model.transformer.positional_encoding import PositionalEncoding
|
||||
from matanyone2.utils.device import safe_autocast
|
||||
|
||||
|
||||
# @torch.jit.script
|
||||
def _weighted_pooling(masks: torch.Tensor, value: torch.Tensor,
|
||||
logits: torch.Tensor) -> (torch.Tensor, torch.Tensor):
|
||||
# value: B*num_objects*H*W*value_dim
|
||||
# logits: B*num_objects*H*W*num_summaries
|
||||
# masks: B*num_objects*H*W*num_summaries: 1 if allowed
|
||||
weights = logits.sigmoid() * masks
|
||||
# B*num_objects*num_summaries*value_dim
|
||||
sums = torch.einsum('bkhwq,bkhwc->bkqc', weights, value)
|
||||
# B*num_objects*H*W*num_summaries -> B*num_objects*num_summaries*1
|
||||
area = weights.flatten(start_dim=2, end_dim=3).sum(2).unsqueeze(-1)
|
||||
|
||||
# B*num_objects*num_summaries*value_dim
|
||||
return sums, area
|
||||
|
||||
|
||||
class ObjectSummarizer(nn.Module):
|
||||
def __init__(self, model_cfg: DictConfig):
|
||||
super().__init__()
|
||||
|
||||
this_cfg = model_cfg.object_summarizer
|
||||
self.value_dim = model_cfg.value_dim
|
||||
self.embed_dim = this_cfg.embed_dim
|
||||
self.num_summaries = this_cfg.num_summaries
|
||||
self.add_pe = this_cfg.add_pe
|
||||
self.pixel_pe_scale = model_cfg.pixel_pe_scale
|
||||
self.pixel_pe_temperature = model_cfg.pixel_pe_temperature
|
||||
|
||||
if self.add_pe:
|
||||
self.pos_enc = PositionalEncoding(self.embed_dim,
|
||||
scale=self.pixel_pe_scale,
|
||||
temperature=self.pixel_pe_temperature)
|
||||
|
||||
self.input_proj = nn.Linear(self.value_dim, self.embed_dim)
|
||||
self.feature_pred = nn.Sequential(
|
||||
nn.Linear(self.embed_dim, self.embed_dim),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Linear(self.embed_dim, self.embed_dim),
|
||||
)
|
||||
self.weights_pred = nn.Sequential(
|
||||
nn.Linear(self.embed_dim, self.embed_dim),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Linear(self.embed_dim, self.num_summaries),
|
||||
)
|
||||
|
||||
def forward(self,
|
||||
masks: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
need_weights: bool = False) -> (torch.Tensor, Optional[torch.Tensor]):
|
||||
# masks: B*num_objects*(H0)*(W0)
|
||||
# value: B*num_objects*value_dim*H*W
|
||||
# -> B*num_objects*H*W*value_dim
|
||||
h, w = value.shape[-2:]
|
||||
masks = F.interpolate(masks, size=(h, w), mode='area')
|
||||
masks = masks.unsqueeze(-1)
|
||||
inv_masks = 1 - masks
|
||||
repeated_masks = torch.cat([
|
||||
masks.expand(-1, -1, -1, -1, self.num_summaries // 2),
|
||||
inv_masks.expand(-1, -1, -1, -1, self.num_summaries // 2),
|
||||
],
|
||||
dim=-1)
|
||||
|
||||
value = value.permute(0, 1, 3, 4, 2)
|
||||
value = self.input_proj(value)
|
||||
if self.add_pe:
|
||||
pe = self.pos_enc(value)
|
||||
value = value + pe
|
||||
|
||||
with safe_autocast(enabled=False): # autocast disabled intentionally
|
||||
value = value.float()
|
||||
feature = self.feature_pred(value)
|
||||
logits = self.weights_pred(value)
|
||||
sums, area = _weighted_pooling(repeated_masks, feature, logits)
|
||||
|
||||
summaries = torch.cat([sums, area], dim=-1)
|
||||
|
||||
if need_weights:
|
||||
return summaries, logits
|
||||
else:
|
||||
return summaries, None
|
||||
@@ -0,0 +1,206 @@
|
||||
from typing import Dict, Optional
|
||||
from omegaconf import DictConfig
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from matanyone2.model.group_modules import GConv2d
|
||||
from matanyone2.utils.tensor_utils import aggregate
|
||||
from matanyone2.model.transformer.positional_encoding import PositionalEncoding
|
||||
from matanyone2.model.transformer.transformer_layers import CrossAttention, SelfAttention, FFN, PixelFFN
|
||||
|
||||
|
||||
class QueryTransformerBlock(nn.Module):
|
||||
def __init__(self, model_cfg: DictConfig):
|
||||
super().__init__()
|
||||
|
||||
this_cfg = model_cfg.object_transformer
|
||||
self.embed_dim = this_cfg.embed_dim
|
||||
self.num_heads = this_cfg.num_heads
|
||||
self.num_queries = this_cfg.num_queries
|
||||
self.ff_dim = this_cfg.ff_dim
|
||||
|
||||
self.read_from_pixel = CrossAttention(self.embed_dim,
|
||||
self.num_heads,
|
||||
add_pe_to_qkv=this_cfg.read_from_pixel.add_pe_to_qkv)
|
||||
self.self_attn = SelfAttention(self.embed_dim,
|
||||
self.num_heads,
|
||||
add_pe_to_qkv=this_cfg.query_self_attention.add_pe_to_qkv)
|
||||
self.ffn = FFN(self.embed_dim, self.ff_dim)
|
||||
self.read_from_query = CrossAttention(self.embed_dim,
|
||||
self.num_heads,
|
||||
add_pe_to_qkv=this_cfg.read_from_query.add_pe_to_qkv,
|
||||
norm=this_cfg.read_from_query.output_norm)
|
||||
self.pixel_ffn = PixelFFN(self.embed_dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
pixel: torch.Tensor,
|
||||
query_pe: torch.Tensor,
|
||||
pixel_pe: torch.Tensor,
|
||||
attn_mask: torch.Tensor,
|
||||
need_weights: bool = False) -> (torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor):
|
||||
# x: (bs*num_objects)*num_queries*embed_dim
|
||||
# pixel: bs*num_objects*C*H*W
|
||||
# query_pe: (bs*num_objects)*num_queries*embed_dim
|
||||
# pixel_pe: (bs*num_objects)*(H*W)*C
|
||||
# attn_mask: (bs*num_objects*num_heads)*num_queries*(H*W)
|
||||
|
||||
# bs*num_objects*C*H*W -> (bs*num_objects)*(H*W)*C
|
||||
pixel_flat = pixel.flatten(3, 4).flatten(0, 1).transpose(1, 2).contiguous()
|
||||
x, q_weights = self.read_from_pixel(x,
|
||||
pixel_flat,
|
||||
query_pe,
|
||||
pixel_pe,
|
||||
attn_mask=attn_mask,
|
||||
need_weights=need_weights)
|
||||
x = self.self_attn(x, query_pe)
|
||||
x = self.ffn(x)
|
||||
|
||||
pixel_flat, p_weights = self.read_from_query(pixel_flat,
|
||||
x,
|
||||
pixel_pe,
|
||||
query_pe,
|
||||
need_weights=need_weights)
|
||||
pixel = self.pixel_ffn(pixel, pixel_flat)
|
||||
|
||||
if need_weights:
|
||||
bs, num_objects, _, h, w = pixel.shape
|
||||
q_weights = q_weights.view(bs, num_objects, self.num_heads, self.num_queries, h, w)
|
||||
p_weights = p_weights.transpose(2, 3).view(bs, num_objects, self.num_heads,
|
||||
self.num_queries, h, w)
|
||||
|
||||
return x, pixel, q_weights, p_weights
|
||||
|
||||
|
||||
class QueryTransformer(nn.Module):
|
||||
def __init__(self, model_cfg: DictConfig):
|
||||
super().__init__()
|
||||
|
||||
this_cfg = model_cfg.object_transformer
|
||||
self.value_dim = model_cfg.value_dim
|
||||
self.embed_dim = this_cfg.embed_dim
|
||||
self.num_heads = this_cfg.num_heads
|
||||
self.num_queries = this_cfg.num_queries
|
||||
|
||||
# query initialization and embedding
|
||||
self.query_init = nn.Embedding(self.num_queries, self.embed_dim)
|
||||
self.query_emb = nn.Embedding(self.num_queries, self.embed_dim)
|
||||
|
||||
# projection from object summaries to query initialization and embedding
|
||||
self.summary_to_query_init = nn.Linear(self.embed_dim, self.embed_dim)
|
||||
self.summary_to_query_emb = nn.Linear(self.embed_dim, self.embed_dim)
|
||||
|
||||
self.pixel_pe_scale = model_cfg.pixel_pe_scale
|
||||
self.pixel_pe_temperature = model_cfg.pixel_pe_temperature
|
||||
self.pixel_init_proj = GConv2d(self.embed_dim, self.embed_dim, kernel_size=1)
|
||||
self.pixel_emb_proj = GConv2d(self.embed_dim, self.embed_dim, kernel_size=1)
|
||||
self.spatial_pe = PositionalEncoding(self.embed_dim,
|
||||
scale=self.pixel_pe_scale,
|
||||
temperature=self.pixel_pe_temperature,
|
||||
channel_last=False,
|
||||
transpose_output=True)
|
||||
|
||||
# transformer blocks
|
||||
self.num_blocks = this_cfg.num_blocks
|
||||
self.blocks = nn.ModuleList(
|
||||
QueryTransformerBlock(model_cfg) for _ in range(self.num_blocks))
|
||||
self.mask_pred = nn.ModuleList(
|
||||
nn.Sequential(nn.ReLU(), GConv2d(self.embed_dim, 1, kernel_size=1))
|
||||
for _ in range(self.num_blocks + 1))
|
||||
|
||||
self.act = nn.ReLU(inplace=True)
|
||||
|
||||
def forward(self,
|
||||
pixel: torch.Tensor,
|
||||
obj_summaries: torch.Tensor,
|
||||
selector: Optional[torch.Tensor] = None,
|
||||
need_weights: bool = False,
|
||||
seg_pass=False) -> (torch.Tensor, Dict[str, torch.Tensor]):
|
||||
# pixel: B*num_objects*embed_dim*H*W
|
||||
# obj_summaries: B*num_objects*T*num_queries*embed_dim
|
||||
T = obj_summaries.shape[2]
|
||||
bs, num_objects, _, H, W = pixel.shape
|
||||
|
||||
# normalize object values
|
||||
# the last channel is the cumulative area of the object
|
||||
obj_summaries = obj_summaries.view(bs * num_objects, T, self.num_queries,
|
||||
self.embed_dim + 1)
|
||||
# sum over time
|
||||
# during inference, T=1 as we already did streaming average in memory_manager
|
||||
obj_sums = obj_summaries[:, :, :, :-1].sum(dim=1)
|
||||
obj_area = obj_summaries[:, :, :, -1:].sum(dim=1)
|
||||
obj_values = obj_sums / (obj_area + 1e-4)
|
||||
obj_init = self.summary_to_query_init(obj_values)
|
||||
obj_emb = self.summary_to_query_emb(obj_values)
|
||||
|
||||
# positional embeddings for object queries
|
||||
query = self.query_init.weight.unsqueeze(0).expand(bs * num_objects, -1, -1) + obj_init
|
||||
query_emb = self.query_emb.weight.unsqueeze(0).expand(bs * num_objects, -1, -1) + obj_emb
|
||||
|
||||
# positional embeddings for pixel features
|
||||
pixel_init = self.pixel_init_proj(pixel)
|
||||
pixel_emb = self.pixel_emb_proj(pixel)
|
||||
pixel_pe = self.spatial_pe(pixel.flatten(0, 1))
|
||||
pixel_emb = pixel_emb.flatten(3, 4).flatten(0, 1).transpose(1, 2).contiguous()
|
||||
pixel_pe = pixel_pe.flatten(1, 2) + pixel_emb
|
||||
|
||||
pixel = pixel_init
|
||||
|
||||
# run the transformer
|
||||
aux_features = {'logits': []}
|
||||
|
||||
# first aux output
|
||||
aux_logits = self.mask_pred[0](pixel).squeeze(2)
|
||||
attn_mask = self._get_aux_mask(aux_logits, selector, seg_pass=seg_pass)
|
||||
aux_features['logits'].append(aux_logits)
|
||||
for i in range(self.num_blocks):
|
||||
query, pixel, q_weights, p_weights = self.blocks[i](query,
|
||||
pixel,
|
||||
query_emb,
|
||||
pixel_pe,
|
||||
attn_mask,
|
||||
need_weights=need_weights)
|
||||
|
||||
if self.training or i <= self.num_blocks - 1 or need_weights:
|
||||
aux_logits = self.mask_pred[i + 1](pixel).squeeze(2)
|
||||
attn_mask = self._get_aux_mask(aux_logits, selector, seg_pass=seg_pass)
|
||||
aux_features['logits'].append(aux_logits)
|
||||
|
||||
aux_features['q_weights'] = q_weights # last layer only
|
||||
aux_features['p_weights'] = p_weights # last layer only
|
||||
|
||||
if self.training:
|
||||
# no need to save all heads
|
||||
aux_features['attn_mask'] = attn_mask.view(bs, num_objects, self.num_heads,
|
||||
self.num_queries, H, W)[:, :, 0]
|
||||
|
||||
return pixel, aux_features
|
||||
|
||||
def _get_aux_mask(self, logits: torch.Tensor, selector: torch.Tensor, seg_pass=False) -> torch.Tensor:
|
||||
# logits: batch_size*num_objects*H*W
|
||||
# selector: batch_size*num_objects*1*1
|
||||
# returns a mask of shape (batch_size*num_objects*num_heads)*num_queries*(H*W)
|
||||
# where True means the attention is blocked
|
||||
|
||||
if selector is None:
|
||||
prob = logits.sigmoid()
|
||||
else:
|
||||
prob = logits.sigmoid() * selector
|
||||
logits = aggregate(prob, dim=1)
|
||||
|
||||
is_foreground = (logits[:, 1:] >= logits.max(dim=1, keepdim=True)[0])
|
||||
foreground_mask = is_foreground.bool().flatten(start_dim=2)
|
||||
inv_foreground_mask = ~foreground_mask
|
||||
inv_background_mask = foreground_mask
|
||||
|
||||
aux_foreground_mask = inv_foreground_mask.unsqueeze(2).unsqueeze(2).repeat(
|
||||
1, 1, self.num_heads, self.num_queries // 2, 1).flatten(start_dim=0, end_dim=2)
|
||||
aux_background_mask = inv_background_mask.unsqueeze(2).unsqueeze(2).repeat(
|
||||
1, 1, self.num_heads, self.num_queries // 2, 1).flatten(start_dim=0, end_dim=2)
|
||||
|
||||
aux_mask = torch.cat([aux_foreground_mask, aux_background_mask], dim=1)
|
||||
|
||||
aux_mask[torch.where(aux_mask.sum(-1) == aux_mask.shape[-1])] = False
|
||||
|
||||
return aux_mask
|
||||
@@ -0,0 +1,110 @@
|
||||
# Reference:
|
||||
# https://github.com/facebookresearch/Mask2Former/blob/main/mask2former/modeling/transformer_decoder/position_encoding.py
|
||||
# https://github.com/tatp22/multidim-positional-encoding/blob/master/positional_encodings/torch_encodings.py
|
||||
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
from matanyone2.utils.device import get_default_device
|
||||
|
||||
|
||||
def get_emb(sin_inp: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Gets a base embedding for one dimension with sin and cos intertwined
|
||||
"""
|
||||
emb = torch.stack((sin_inp.sin(), sin_inp.cos()), dim=-1)
|
||||
return torch.flatten(emb, -2, -1)
|
||||
|
||||
|
||||
class PositionalEncoding(nn.Module):
|
||||
def __init__(self,
|
||||
dim: int,
|
||||
scale: float = math.pi * 2,
|
||||
temperature: float = 10000,
|
||||
normalize: bool = True,
|
||||
channel_last: bool = True,
|
||||
transpose_output: bool = False):
|
||||
super().__init__()
|
||||
dim = int(np.ceil(dim / 4) * 2)
|
||||
self.dim = dim
|
||||
inv_freq = 1.0 / (temperature**(torch.arange(0, dim, 2).float() / dim))
|
||||
self.register_buffer("inv_freq", inv_freq)
|
||||
self.normalize = normalize
|
||||
self.scale = scale
|
||||
self.eps = 1e-6
|
||||
self.channel_last = channel_last
|
||||
self.transpose_output = transpose_output
|
||||
|
||||
self.cached_penc = None # the cache is irrespective of the number of objects
|
||||
|
||||
def forward(self, tensor: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
:param tensor: A 4/5d tensor of size
|
||||
channel_last=True: (batch_size, h, w, c) or (batch_size, k, h, w, c)
|
||||
channel_last=False: (batch_size, c, h, w) or (batch_size, k, c, h, w)
|
||||
:return: positional encoding tensor that has the same shape as the input if the input is 4d
|
||||
if the input is 5d, the output is broadcastable along the k-dimension
|
||||
"""
|
||||
if len(tensor.shape) != 4 and len(tensor.shape) != 5:
|
||||
raise RuntimeError(f'The input tensor has to be 4/5d, got {tensor.shape}!')
|
||||
|
||||
if len(tensor.shape) == 5:
|
||||
# take a sample from the k dimension
|
||||
num_objects = tensor.shape[1]
|
||||
tensor = tensor[:, 0]
|
||||
else:
|
||||
num_objects = None
|
||||
|
||||
if self.channel_last:
|
||||
batch_size, h, w, c = tensor.shape
|
||||
else:
|
||||
batch_size, c, h, w = tensor.shape
|
||||
|
||||
if self.cached_penc is not None and self.cached_penc.shape == tensor.shape:
|
||||
if num_objects is None:
|
||||
return self.cached_penc
|
||||
else:
|
||||
return self.cached_penc.unsqueeze(1)
|
||||
|
||||
self.cached_penc = None
|
||||
|
||||
pos_y = torch.arange(h, device=tensor.device, dtype=self.inv_freq.dtype)
|
||||
pos_x = torch.arange(w, device=tensor.device, dtype=self.inv_freq.dtype)
|
||||
if self.normalize:
|
||||
pos_y = pos_y / (pos_y[-1] + self.eps) * self.scale
|
||||
pos_x = pos_x / (pos_x[-1] + self.eps) * self.scale
|
||||
|
||||
sin_inp_y = torch.einsum("i,j->ij", pos_y, self.inv_freq)
|
||||
sin_inp_x = torch.einsum("i,j->ij", pos_x, self.inv_freq)
|
||||
emb_y = get_emb(sin_inp_y).unsqueeze(1)
|
||||
emb_x = get_emb(sin_inp_x)
|
||||
|
||||
emb = torch.zeros((h, w, self.dim * 2), device=tensor.device, dtype=tensor.dtype)
|
||||
emb[:, :, :self.dim] = emb_x
|
||||
emb[:, :, self.dim:] = emb_y
|
||||
|
||||
if not self.channel_last and self.transpose_output:
|
||||
# cancelled out
|
||||
pass
|
||||
elif (not self.channel_last) or (self.transpose_output):
|
||||
emb = emb.permute(2, 0, 1)
|
||||
|
||||
self.cached_penc = emb.unsqueeze(0).repeat(batch_size, 1, 1, 1)
|
||||
if num_objects is None:
|
||||
return self.cached_penc
|
||||
else:
|
||||
return self.cached_penc.unsqueeze(1)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
device = get_default_device()
|
||||
pe = PositionalEncoding(8).to(device)
|
||||
input = torch.ones((1, 8, 8, 8), device=device)
|
||||
output = pe(input)
|
||||
# print(output)
|
||||
print(output[0, :, 0, 0])
|
||||
print(output[0, :, 0, 5])
|
||||
print(output[0, 0, :, 0])
|
||||
print(output[0, 0, 0, :])
|
||||
@@ -0,0 +1,161 @@
|
||||
# Modified from PyTorch nn.Transformer
|
||||
|
||||
from typing import List, Callable
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from matanyone2.model.channel_attn import CAResBlock
|
||||
|
||||
|
||||
class SelfAttention(nn.Module):
|
||||
def __init__(self,
|
||||
dim: int,
|
||||
nhead: int,
|
||||
dropout: float = 0.0,
|
||||
batch_first: bool = True,
|
||||
add_pe_to_qkv: List[bool] = [True, True, False]):
|
||||
super().__init__()
|
||||
self.self_attn = nn.MultiheadAttention(dim, nhead, dropout=dropout, batch_first=batch_first)
|
||||
self.norm = nn.LayerNorm(dim)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.add_pe_to_qkv = add_pe_to_qkv
|
||||
|
||||
def forward(self,
|
||||
x: torch.Tensor,
|
||||
pe: torch.Tensor,
|
||||
attn_mask: bool = None,
|
||||
key_padding_mask: bool = None) -> torch.Tensor:
|
||||
x = self.norm(x)
|
||||
if any(self.add_pe_to_qkv):
|
||||
x_with_pe = x + pe
|
||||
q = x_with_pe if self.add_pe_to_qkv[0] else x
|
||||
k = x_with_pe if self.add_pe_to_qkv[1] else x
|
||||
v = x_with_pe if self.add_pe_to_qkv[2] else x
|
||||
else:
|
||||
q = k = v = x
|
||||
|
||||
r = x
|
||||
x = self.self_attn(q, k, v, attn_mask=attn_mask, key_padding_mask=key_padding_mask)[0]
|
||||
return r + self.dropout(x)
|
||||
|
||||
|
||||
# https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html#torch.nn.functional.scaled_dot_product_attention
|
||||
class CrossAttention(nn.Module):
|
||||
def __init__(self,
|
||||
dim: int,
|
||||
nhead: int,
|
||||
dropout: float = 0.0,
|
||||
batch_first: bool = True,
|
||||
add_pe_to_qkv: List[bool] = [True, True, False],
|
||||
residual: bool = True,
|
||||
norm: bool = True):
|
||||
super().__init__()
|
||||
self.cross_attn = nn.MultiheadAttention(dim,
|
||||
nhead,
|
||||
dropout=dropout,
|
||||
batch_first=batch_first)
|
||||
if norm:
|
||||
self.norm = nn.LayerNorm(dim)
|
||||
else:
|
||||
self.norm = nn.Identity()
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.add_pe_to_qkv = add_pe_to_qkv
|
||||
self.residual = residual
|
||||
|
||||
def forward(self,
|
||||
x: torch.Tensor,
|
||||
mem: torch.Tensor,
|
||||
x_pe: torch.Tensor,
|
||||
mem_pe: torch.Tensor,
|
||||
attn_mask: bool = None,
|
||||
*,
|
||||
need_weights: bool = False) -> (torch.Tensor, torch.Tensor):
|
||||
x = self.norm(x)
|
||||
if self.add_pe_to_qkv[0]:
|
||||
q = x + x_pe
|
||||
else:
|
||||
q = x
|
||||
|
||||
if any(self.add_pe_to_qkv[1:]):
|
||||
mem_with_pe = mem + mem_pe
|
||||
k = mem_with_pe if self.add_pe_to_qkv[1] else mem
|
||||
v = mem_with_pe if self.add_pe_to_qkv[2] else mem
|
||||
else:
|
||||
k = v = mem
|
||||
r = x
|
||||
x, weights = self.cross_attn(q,
|
||||
k,
|
||||
v,
|
||||
attn_mask=attn_mask,
|
||||
need_weights=need_weights,
|
||||
average_attn_weights=False)
|
||||
|
||||
if self.residual:
|
||||
return r + self.dropout(x), weights
|
||||
else:
|
||||
return self.dropout(x), weights
|
||||
|
||||
|
||||
class FFN(nn.Module):
|
||||
def __init__(self, dim_in: int, dim_ff: int, activation=F.relu):
|
||||
super().__init__()
|
||||
self.linear1 = nn.Linear(dim_in, dim_ff)
|
||||
self.linear2 = nn.Linear(dim_ff, dim_in)
|
||||
self.norm = nn.LayerNorm(dim_in)
|
||||
|
||||
if isinstance(activation, str):
|
||||
self.activation = _get_activation_fn(activation)
|
||||
else:
|
||||
self.activation = activation
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
r = x
|
||||
x = self.norm(x)
|
||||
x = self.linear2(self.activation(self.linear1(x)))
|
||||
x = r + x
|
||||
return x
|
||||
|
||||
|
||||
class PixelFFN(nn.Module):
|
||||
def __init__(self, dim: int):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.conv = CAResBlock(dim, dim)
|
||||
|
||||
def forward(self, pixel: torch.Tensor, pixel_flat: torch.Tensor) -> torch.Tensor:
|
||||
# pixel: batch_size * num_objects * dim * H * W
|
||||
# pixel_flat: (batch_size*num_objects) * (H*W) * dim
|
||||
bs, num_objects, _, h, w = pixel.shape
|
||||
pixel_flat = pixel_flat.view(bs * num_objects, h, w, self.dim)
|
||||
pixel_flat = pixel_flat.permute(0, 3, 1, 2).contiguous()
|
||||
|
||||
x = self.conv(pixel_flat)
|
||||
x = x.view(bs, num_objects, self.dim, h, w)
|
||||
return x
|
||||
|
||||
|
||||
class OutputFFN(nn.Module):
|
||||
def __init__(self, dim_in: int, dim_out: int, activation=F.relu):
|
||||
super().__init__()
|
||||
self.linear1 = nn.Linear(dim_in, dim_out)
|
||||
self.linear2 = nn.Linear(dim_out, dim_out)
|
||||
|
||||
if isinstance(activation, str):
|
||||
self.activation = _get_activation_fn(activation)
|
||||
else:
|
||||
self.activation = activation
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.linear2(self.activation(self.linear1(x)))
|
||||
return x
|
||||
|
||||
|
||||
def _get_activation_fn(activation: str) -> Callable[[Tensor], Tensor]:
|
||||
if activation == "relu":
|
||||
return F.relu
|
||||
elif activation == "gelu":
|
||||
return F.gelu
|
||||
|
||||
raise RuntimeError("activation should be relu/gelu, not {}".format(activation))
|
||||
@@ -0,0 +1,107 @@
|
||||
import math
|
||||
import torch
|
||||
from typing import Optional, Union, Tuple
|
||||
|
||||
|
||||
# @torch.jit.script
|
||||
def get_similarity(mk: torch.Tensor,
|
||||
ms: torch.Tensor,
|
||||
qk: torch.Tensor,
|
||||
qe: torch.Tensor,
|
||||
add_batch_dim: bool = False,
|
||||
uncert_mask = None) -> torch.Tensor:
|
||||
# used for training/inference and memory reading/memory potentiation
|
||||
# mk: B x CK x [N] - Memory keys
|
||||
# ms: B x 1 x [N] - Memory shrinkage
|
||||
# qk: B x CK x [HW/P] - Query keys
|
||||
# qe: B x CK x [HW/P] - Query selection
|
||||
# Dimensions in [] are flattened
|
||||
# Return: B*N*HW
|
||||
if add_batch_dim:
|
||||
mk, ms = mk.unsqueeze(0), ms.unsqueeze(0)
|
||||
qk, qe = qk.unsqueeze(0), qe.unsqueeze(0)
|
||||
|
||||
CK = mk.shape[1]
|
||||
|
||||
mk = mk.flatten(start_dim=2)
|
||||
ms = ms.flatten(start_dim=1).unsqueeze(2) if ms is not None else None
|
||||
qk = qk.flatten(start_dim=2)
|
||||
qe = qe.flatten(start_dim=2) if qe is not None else None
|
||||
|
||||
# query token selection based on temporal sparsity
|
||||
if uncert_mask is not None:
|
||||
uncert_mask = uncert_mask.flatten(start_dim=2)
|
||||
uncert_mask = uncert_mask.expand(-1, 64, -1)
|
||||
qk = qk * uncert_mask
|
||||
qe = qe * uncert_mask
|
||||
|
||||
if qe is not None:
|
||||
# See XMem's appendix for derivation
|
||||
mk = mk.transpose(1, 2)
|
||||
a_sq = (mk.pow(2) @ qe)
|
||||
two_ab = 2 * (mk @ (qk * qe))
|
||||
b_sq = (qe * qk.pow(2)).sum(1, keepdim=True)
|
||||
similarity = (-a_sq + two_ab - b_sq)
|
||||
else:
|
||||
# similar to STCN if we don't have the selection term
|
||||
a_sq = mk.pow(2).sum(1).unsqueeze(2)
|
||||
two_ab = 2 * (mk.transpose(1, 2) @ qk)
|
||||
similarity = (-a_sq + two_ab)
|
||||
|
||||
if ms is not None:
|
||||
similarity = similarity * ms / math.sqrt(CK) # B*N*HW
|
||||
else:
|
||||
similarity = similarity / math.sqrt(CK) # B*N*HW
|
||||
|
||||
return similarity
|
||||
|
||||
|
||||
def do_softmax(
|
||||
similarity: torch.Tensor,
|
||||
top_k: Optional[int] = None,
|
||||
inplace: bool = False,
|
||||
return_usage: bool = False) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
|
||||
# normalize similarity with top-k softmax
|
||||
# similarity: B x N x [HW/P]
|
||||
# use inplace with care
|
||||
if top_k is not None:
|
||||
values, indices = torch.topk(similarity, k=top_k, dim=1)
|
||||
|
||||
x_exp = values.exp_()
|
||||
x_exp /= torch.sum(x_exp, dim=1, keepdim=True)
|
||||
if inplace:
|
||||
similarity.zero_().scatter_(1, indices, x_exp) # B*N*HW
|
||||
affinity = similarity
|
||||
else:
|
||||
affinity = torch.zeros_like(similarity).scatter_(1, indices, x_exp) # B*N*HW
|
||||
else:
|
||||
maxes = torch.max(similarity, dim=1, keepdim=True)[0]
|
||||
x_exp = torch.exp(similarity - maxes)
|
||||
x_exp_sum = torch.sum(x_exp, dim=1, keepdim=True)
|
||||
affinity = x_exp / x_exp_sum
|
||||
indices = None
|
||||
|
||||
if return_usage:
|
||||
return affinity, affinity.sum(dim=2)
|
||||
|
||||
return affinity
|
||||
|
||||
|
||||
def get_affinity(mk: torch.Tensor, ms: torch.Tensor, qk: torch.Tensor,
|
||||
qe: torch.Tensor, uncert_mask = None) -> torch.Tensor:
|
||||
# shorthand used in training with no top-k
|
||||
similarity = get_similarity(mk, ms, qk, qe, uncert_mask=uncert_mask)
|
||||
affinity = do_softmax(similarity)
|
||||
return affinity
|
||||
|
||||
def readout(affinity: torch.Tensor, mv: torch.Tensor, uncert_mask: torch.Tensor=None) -> torch.Tensor:
|
||||
B, CV, T, H, W = mv.shape
|
||||
|
||||
mo = mv.view(B, CV, T * H * W)
|
||||
mem = torch.bmm(mo, affinity)
|
||||
if uncert_mask is not None:
|
||||
uncert_mask = uncert_mask.flatten(start_dim=2).expand(-1, CV, -1)
|
||||
mem = mem * uncert_mask
|
||||
mem = mem.view(B, CV, H, W)
|
||||
|
||||
return mem
|
||||
@@ -0,0 +1,72 @@
|
||||
import logging
|
||||
|
||||
log = logging.getLogger()
|
||||
|
||||
|
||||
def get_parameter_groups(model, stage_cfg, print_log=False):
|
||||
"""
|
||||
Assign different weight decays and learning rates to different parameters.
|
||||
Returns a parameter group which can be passed to the optimizer.
|
||||
"""
|
||||
weight_decay = stage_cfg.weight_decay
|
||||
embed_weight_decay = stage_cfg.embed_weight_decay
|
||||
backbone_lr_ratio = stage_cfg.backbone_lr_ratio
|
||||
base_lr = stage_cfg.learning_rate
|
||||
|
||||
backbone_params = []
|
||||
embed_params = []
|
||||
other_params = []
|
||||
|
||||
embedding_names = ['summary_pos', 'query_init', 'query_emb', 'obj_pe']
|
||||
embedding_names = [e + '.weight' for e in embedding_names]
|
||||
|
||||
# inspired by detectron2
|
||||
memo = set()
|
||||
for name, param in model.named_parameters():
|
||||
if not param.requires_grad:
|
||||
continue
|
||||
# Avoid duplicating parameters
|
||||
if param in memo:
|
||||
continue
|
||||
memo.add(param)
|
||||
|
||||
if name.startswith('module'):
|
||||
name = name[7:]
|
||||
|
||||
inserted = False
|
||||
if name.startswith('pixel_encoder.'):
|
||||
backbone_params.append(param)
|
||||
inserted = True
|
||||
if print_log:
|
||||
log.info(f'{name} counted as a backbone parameter.')
|
||||
else:
|
||||
for e in embedding_names:
|
||||
if name.endswith(e):
|
||||
embed_params.append(param)
|
||||
inserted = True
|
||||
if print_log:
|
||||
log.info(f'{name} counted as an embedding parameter.')
|
||||
break
|
||||
|
||||
if not inserted:
|
||||
other_params.append(param)
|
||||
|
||||
parameter_groups = [
|
||||
{
|
||||
'params': backbone_params,
|
||||
'lr': base_lr * backbone_lr_ratio,
|
||||
'weight_decay': weight_decay
|
||||
},
|
||||
{
|
||||
'params': embed_params,
|
||||
'lr': base_lr,
|
||||
'weight_decay': embed_weight_decay
|
||||
},
|
||||
{
|
||||
'params': other_params,
|
||||
'lr': base_lr,
|
||||
'weight_decay': weight_decay
|
||||
},
|
||||
]
|
||||
|
||||
return parameter_groups
|
||||
@@ -0,0 +1,179 @@
|
||||
"""
|
||||
resnet.py - A modified ResNet structure
|
||||
We append extra channels to the first conv by some network surgery
|
||||
"""
|
||||
|
||||
from collections import OrderedDict
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils import model_zoo
|
||||
|
||||
|
||||
def load_weights_add_extra_dim(target, source_state, extra_dim=1):
|
||||
new_dict = OrderedDict()
|
||||
|
||||
for k1, v1 in target.state_dict().items():
|
||||
if 'num_batches_tracked' not in k1:
|
||||
if k1 in source_state:
|
||||
tar_v = source_state[k1]
|
||||
|
||||
if v1.shape != tar_v.shape:
|
||||
# Init the new segmentation channel with zeros
|
||||
# print(v1.shape, tar_v.shape)
|
||||
c, _, w, h = v1.shape
|
||||
pads = torch.zeros((c, extra_dim, w, h), device=tar_v.device)
|
||||
nn.init.orthogonal_(pads)
|
||||
tar_v = torch.cat([tar_v, pads], 1)
|
||||
|
||||
new_dict[k1] = tar_v
|
||||
|
||||
target.load_state_dict(new_dict)
|
||||
|
||||
|
||||
model_urls = {
|
||||
'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
|
||||
'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
|
||||
}
|
||||
|
||||
|
||||
def conv3x3(in_planes, out_planes, stride=1, dilation=1):
|
||||
return nn.Conv2d(in_planes,
|
||||
out_planes,
|
||||
kernel_size=3,
|
||||
stride=stride,
|
||||
padding=dilation,
|
||||
dilation=dilation,
|
||||
bias=False)
|
||||
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None, dilation=1):
|
||||
super(BasicBlock, self).__init__()
|
||||
self.conv1 = conv3x3(inplanes, planes, stride=stride, dilation=dilation)
|
||||
self.bn1 = nn.BatchNorm2d(planes)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.conv2 = conv3x3(planes, planes, stride=1, dilation=dilation)
|
||||
self.bn2 = nn.BatchNorm2d(planes)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
residual = self.downsample(x)
|
||||
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class Bottleneck(nn.Module):
|
||||
expansion = 4
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None, dilation=1):
|
||||
super(Bottleneck, self).__init__()
|
||||
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(planes)
|
||||
self.conv2 = nn.Conv2d(planes,
|
||||
planes,
|
||||
kernel_size=3,
|
||||
stride=stride,
|
||||
dilation=dilation,
|
||||
padding=dilation,
|
||||
bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(planes)
|
||||
self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes * 4)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
residual = self.downsample(x)
|
||||
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ResNet(nn.Module):
|
||||
def __init__(self, block, layers=(3, 4, 23, 3), extra_dim=0):
|
||||
self.inplanes = 64
|
||||
super(ResNet, self).__init__()
|
||||
self.conv1 = nn.Conv2d(3 + extra_dim, 64, kernel_size=7, stride=2, padding=3, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(64)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
|
||||
self.layer1 = self._make_layer(block, 64, layers[0])
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
|
||||
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
|
||||
m.weight.data.normal_(0, math.sqrt(2. / n))
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
m.weight.data.fill_(1)
|
||||
m.bias.data.zero_()
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1, dilation=1):
|
||||
downsample = None
|
||||
if stride != 1 or self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
nn.Conv2d(self.inplanes,
|
||||
planes * block.expansion,
|
||||
kernel_size=1,
|
||||
stride=stride,
|
||||
bias=False),
|
||||
nn.BatchNorm2d(planes * block.expansion),
|
||||
)
|
||||
|
||||
layers = [block(self.inplanes, planes, stride, downsample)]
|
||||
self.inplanes = planes * block.expansion
|
||||
for i in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes, dilation=dilation))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
|
||||
def resnet18(pretrained=True, extra_dim=0):
|
||||
model = ResNet(BasicBlock, [2, 2, 2, 2], extra_dim)
|
||||
if pretrained:
|
||||
load_weights_add_extra_dim(model, model_zoo.load_url(model_urls['resnet18']), extra_dim)
|
||||
return model
|
||||
|
||||
|
||||
def resnet50(pretrained=True, extra_dim=0):
|
||||
model = ResNet(Bottleneck, [3, 4, 6, 3], extra_dim)
|
||||
if pretrained:
|
||||
load_weights_add_extra_dim(model, model_zoo.load_url(model_urls['resnet50']), extra_dim)
|
||||
return model
|
||||
Reference in New Issue
Block a user