|
|
|
|
|
|
|
try: |
|
from timm.layers import resample_abs_pos_embed |
|
except ImportError as err: |
|
print("ImportError: {0}".format(err)) |
|
import torch |
|
import torch.nn as nn |
|
from torch.utils.checkpoint import checkpoint |
|
|
|
|
|
def make_vit_b16_backbone( |
|
model, |
|
encoder_feature_dims, |
|
encoder_feature_layer_ids, |
|
vit_features, |
|
start_index=1, |
|
use_grad_checkpointing=False, |
|
) -> nn.Module: |
|
"""Make a ViTb16 backbone for the DPT model.""" |
|
if use_grad_checkpointing: |
|
model.set_grad_checkpointing() |
|
|
|
vit_model = nn.Module() |
|
vit_model.hooks = encoder_feature_layer_ids |
|
vit_model.model = model |
|
vit_model.features = encoder_feature_dims |
|
vit_model.vit_features = vit_features |
|
vit_model.model.start_index = start_index |
|
vit_model.model.patch_size = vit_model.model.patch_embed.patch_size |
|
vit_model.model.is_vit = True |
|
vit_model.model.forward = vit_model.model.forward_features |
|
|
|
return vit_model |
|
|
|
|
|
def forward_features_eva_fixed(self, x): |
|
"""Encode features.""" |
|
x = self.patch_embed(x) |
|
x, rot_pos_embed = self._pos_embed(x) |
|
for blk in self.blocks: |
|
if self.grad_checkpointing: |
|
x = checkpoint(blk, x, rot_pos_embed) |
|
else: |
|
x = blk(x, rot_pos_embed) |
|
x = self.norm(x) |
|
return x |
|
|
|
|
|
def resize_vit(model: nn.Module, img_size) -> nn.Module: |
|
"""Resample the ViT module to the given size.""" |
|
patch_size = model.patch_embed.patch_size |
|
model.patch_embed.img_size = img_size |
|
grid_size = tuple([s // p for s, p in zip(img_size, patch_size)]) |
|
model.patch_embed.grid_size = grid_size |
|
|
|
pos_embed = resample_abs_pos_embed( |
|
model.pos_embed, |
|
grid_size, |
|
num_prefix_tokens=( |
|
0 if getattr(model, "no_embed_class", False) else model.num_prefix_tokens |
|
), |
|
) |
|
model.pos_embed = torch.nn.Parameter(pos_embed) |
|
|
|
return model |
|
|
|
|
|
def resize_patch_embed(model: nn.Module, new_patch_size=(16, 16)) -> nn.Module: |
|
"""Resample the ViT patch size to the given one.""" |
|
|
|
if hasattr(model, "patch_embed"): |
|
old_patch_size = model.patch_embed.patch_size |
|
|
|
if ( |
|
new_patch_size[0] != old_patch_size[0] |
|
or new_patch_size[1] != old_patch_size[1] |
|
): |
|
patch_embed_proj = model.patch_embed.proj.weight |
|
patch_embed_proj_bias = model.patch_embed.proj.bias |
|
use_bias = True if patch_embed_proj_bias is not None else False |
|
_, _, h, w = patch_embed_proj.shape |
|
|
|
new_patch_embed_proj = torch.nn.functional.interpolate( |
|
patch_embed_proj, |
|
size=[new_patch_size[0], new_patch_size[1]], |
|
mode="bicubic", |
|
align_corners=False, |
|
) |
|
new_patch_embed_proj = ( |
|
new_patch_embed_proj * (h / new_patch_size[0]) * (w / new_patch_size[1]) |
|
) |
|
|
|
model.patch_embed.proj = nn.Conv2d( |
|
in_channels=model.patch_embed.proj.in_channels, |
|
out_channels=model.patch_embed.proj.out_channels, |
|
kernel_size=new_patch_size, |
|
stride=new_patch_size, |
|
bias=use_bias, |
|
) |
|
|
|
if use_bias: |
|
model.patch_embed.proj.bias = patch_embed_proj_bias |
|
|
|
model.patch_embed.proj.weight = torch.nn.Parameter(new_patch_embed_proj) |
|
|
|
model.patch_size = new_patch_size |
|
model.patch_embed.patch_size = new_patch_size |
|
model.patch_embed.img_size = ( |
|
int( |
|
model.patch_embed.img_size[0] |
|
* new_patch_size[0] |
|
/ old_patch_size[0] |
|
), |
|
int( |
|
model.patch_embed.img_size[1] |
|
* new_patch_size[1] |
|
/ old_patch_size[1] |
|
), |
|
) |
|
|
|
return model |
|
|
