release infer and demo
This commit is contained in:
pq-yang
@@ -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