|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.nn.functional as F |
|
|
from typing import List, Tuple |
|
|
from torch import Tensor |
|
|
from open_clip import get_tokenizer, create_model_from_pretrained |
|
|
import torchvision.transforms as T |
|
|
from .utils import imagenet_templates |
|
|
|
|
|
OPENAI_NORMALIZE = T.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)) |
|
|
|
|
|
class MaskClip(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
clip_model="ViT-B-16", |
|
|
pretrained="laion2b_s34b_b88k", |
|
|
patch_size=16, |
|
|
img_size=(224, 224), |
|
|
in_channels=768, |
|
|
text_channels=512, |
|
|
): |
|
|
super(MaskClip, self).__init__() |
|
|
|
|
|
self.patch_size = patch_size |
|
|
self.img_size = img_size |
|
|
model, _ = create_model_from_pretrained(clip_model, pretrained=pretrained) |
|
|
model.eval() |
|
|
self.clip_T = OPENAI_NORMALIZE |
|
|
self.hook_features = {} |
|
|
self.backbone = model |
|
|
def hook_fn_forward(module, input, output): |
|
|
self.hook_features["v"] = output |
|
|
self.backbone.visual.transformer.resblocks[-2].register_forward_hook(hook_fn_forward) |
|
|
self._positional_embd = nn.Parameter(self.backbone.visual.positional_embedding.data.clone()) |
|
|
self.proj = nn.Conv2d(in_channels, text_channels, 1, bias=False) |
|
|
self.proj.weight = nn.Parameter(model.visual.proj.t()[:, :, None, None]) |
|
|
self.tokenizer = get_tokenizer(clip_model) |
|
|
|
|
|
@torch.no_grad() |
|
|
def extract_feat(self, inputs: Tensor) -> Tuple[Tensor]: |
|
|
"""Extract features from images.""" |
|
|
pos_embed = self.backbone.visual.positional_embedding |
|
|
|
|
|
B, C, H, W = inputs.shape |
|
|
hw_shape = (H // self.patch_size, W // self.patch_size) |
|
|
x_len, pos_len = hw_shape[0]*hw_shape[1], pos_embed.shape[0] |
|
|
|
|
|
if x_len != pos_len: |
|
|
if pos_len == (self.img_size[0] // self.patch_size) * (self.img_size[1] // self.patch_size) + 1: |
|
|
pos_h = self.img_size[0] // self.patch_size |
|
|
pos_w = self.img_size[1] // self.patch_size |
|
|
else: |
|
|
raise ValueError( |
|
|
'{}, {}'.format(x_len, pos_len)) |
|
|
|
|
|
self.backbone.visual.positional_embedding.data = self.resize_pos_embed( |
|
|
self._positional_embd[None], hw_shape, (pos_h, pos_w), 'bicubic')[0] |
|
|
|
|
|
_ = self.backbone(inputs) |
|
|
v = self.hook_features["v"] |
|
|
v = self.extract_v(v, self.backbone.visual.transformer.resblocks[-1]).permute(1, 0, 2) |
|
|
v = self.backbone.visual.ln_post(v) |
|
|
|
|
|
v = v.permute(1, 0, 2)[:, 1:] |
|
|
v = v.reshape(B, hw_shape[0], hw_shape[1], -1).permute(0, 3, 1, 2).contiguous() |
|
|
|
|
|
self.backbone.visual.positional_embedding.data = self._positional_embd |
|
|
return v |
|
|
|
|
|
@torch.no_grad() |
|
|
def extract_v(self, x, block): |
|
|
y = block.ln_1(x) |
|
|
y = torch.nn.functional.linear(y, block.attn.in_proj_weight, block.attn.in_proj_bias) |
|
|
B, N, C = y.shape |
|
|
y = y.view(B, N, 3, C // 3).permute(2, 0, 1, 3).reshape(3 * B, N, C // 3) |
|
|
y = F.linear(y, block.attn.out_proj.weight, block.attn.out_proj.bias) |
|
|
q, k, v = y.tensor_split(3, dim=0) |
|
|
v += x |
|
|
v += block.mlp(block.ln_2(v)) |
|
|
return v |
|
|
|
|
|
|
|
|
@staticmethod |
|
|
def resize_pos_embed(pos_embed, input_shpae, pos_shape, mode): |
|
|
"""Resize pos_embed weights. |
|
|
|
|
|
Resize pos_embed using bicubic interpolate method. |
|
|
Args: |
|
|
pos_embed (torch.Tensor): Position embedding weights. |
|
|
input_shpae (tuple): Tuple for (downsampled input image height, |
|
|
downsampled input image width). |
|
|
pos_shape (tuple): The resolution of downsampled origin training |
|
|
image. |
|
|
mode (str): Algorithm used for upsampling: |
|
|
``'nearest'`` | ``'linear'`` | ``'bilinear'`` | ``'bicubic'`` | |
|
|
``'trilinear'``. Default: ``'nearest'`` |
|
|
Return: |
|
|
torch.Tensor: The resized pos_embed of shape [B, L_new, C] |
|
|
""" |
|
|
assert pos_embed.ndim == 3, 'shape of pos_embed must be [B, L, C]' |
|
|
pos_h, pos_w = pos_shape |
|
|
cls_token_weight = pos_embed[:, 0] |
|
|
pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w):] |
|
|
pos_embed_weight = pos_embed_weight.reshape( |
|
|
1, pos_h, pos_w, pos_embed.shape[2]).permute(0, 3, 1, 2) |
|
|
pos_embed_weight = F.interpolate( |
|
|
pos_embed_weight, size=input_shpae, align_corners=False, mode=mode) |
|
|
cls_token_weight = cls_token_weight.unsqueeze(1) |
|
|
pos_embed_weight = torch.flatten(pos_embed_weight, 2).transpose(1, 2) |
|
|
pos_embed = torch.cat((cls_token_weight, pos_embed_weight), dim=1) |
|
|
return pos_embed |
|
|
|
|
|
@torch.no_grad() |
|
|
def decode_head(self, x: Tensor) -> Tensor: |
|
|
feat = self.proj(x) |
|
|
|
|
|
return feat |
|
|
|
|
|
|
|
|
@torch.no_grad() |
|
|
def forward(self, inputs: Tensor) -> Tensor: |
|
|
"""Encode images with backbone and decode into a semantic segmentation |
|
|
map of the same size as input.""" |
|
|
inputs = self.clip_T(inputs) |
|
|
x = self.extract_feat(inputs) |
|
|
feats = self.decode_head(x) |
|
|
return feats |
|
|
|
|
|
|
|
|
@torch.no_grad() |
|
|
def get_classifier(self, classnames:List[str]) -> Tensor: |
|
|
aug_embeddings = torch.stack([self._embed_label(label) for label in classnames]) |
|
|
aug_embeddings = aug_embeddings / aug_embeddings.norm(dim=-1, keepdim=True) |
|
|
return aug_embeddings.squeeze(1) |
|
|
|
|
|
|
|
|
@torch.no_grad() |
|
|
def _embed_label(self, label: str) -> Tensor: |
|
|
"""Encode label name into a single vector.""" |
|
|
all_prompts = [self.tokenizer(template.format(label)) for template in imagenet_templates] |
|
|
all_prompts = torch.cat(all_prompts) |
|
|
all_prompts = all_prompts.to(self.backbone.visual.positional_embedding.device) |
|
|
out = self.backbone.encode_text(all_prompts) |
|
|
out /= out.norm(dim=-1, keepdim=True) |
|
|
out = out.mean(dim=0) |
|
|
return out |
|
|
|
