|
import timm |
|
import torch |
|
import torch.nn as nn |
|
import numpy as np |
|
|
|
from .utils import activations, get_activation, Transpose |
|
|
|
|
|
def forward_levit(pretrained, x): |
|
pretrained.model.forward_features(x) |
|
|
|
layer_1 = pretrained.activations["1"] |
|
layer_2 = pretrained.activations["2"] |
|
layer_3 = pretrained.activations["3"] |
|
|
|
layer_1 = pretrained.act_postprocess1(layer_1) |
|
layer_2 = pretrained.act_postprocess2(layer_2) |
|
layer_3 = pretrained.act_postprocess3(layer_3) |
|
|
|
return layer_1, layer_2, layer_3 |
|
|
|
|
|
def _make_levit_backbone( |
|
model, |
|
hooks=[3, 11, 21], |
|
patch_grid=[14, 14] |
|
): |
|
pretrained = nn.Module() |
|
|
|
pretrained.model = model |
|
pretrained.model.blocks[hooks[0]].register_forward_hook(get_activation("1")) |
|
pretrained.model.blocks[hooks[1]].register_forward_hook(get_activation("2")) |
|
pretrained.model.blocks[hooks[2]].register_forward_hook(get_activation("3")) |
|
|
|
pretrained.activations = activations |
|
|
|
patch_grid_size = np.array(patch_grid, dtype=int) |
|
|
|
pretrained.act_postprocess1 = nn.Sequential( |
|
Transpose(1, 2), |
|
nn.Unflatten(2, torch.Size(patch_grid_size.tolist())) |
|
) |
|
pretrained.act_postprocess2 = nn.Sequential( |
|
Transpose(1, 2), |
|
nn.Unflatten(2, torch.Size((np.ceil(patch_grid_size / 2).astype(int)).tolist())) |
|
) |
|
pretrained.act_postprocess3 = nn.Sequential( |
|
Transpose(1, 2), |
|
nn.Unflatten(2, torch.Size((np.ceil(patch_grid_size / 4).astype(int)).tolist())) |
|
) |
|
|
|
return pretrained |
|
|
|
|
|
class ConvTransposeNorm(nn.Sequential): |
|
""" |
|
Modification of |
|
https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/levit.py: ConvNorm |
|
such that ConvTranspose2d is used instead of Conv2d. |
|
""" |
|
|
|
def __init__( |
|
self, in_chs, out_chs, kernel_size=1, stride=1, pad=0, dilation=1, |
|
groups=1, bn_weight_init=1): |
|
super().__init__() |
|
self.add_module('c', |
|
nn.ConvTranspose2d(in_chs, out_chs, kernel_size, stride, pad, dilation, groups, bias=False)) |
|
self.add_module('bn', nn.BatchNorm2d(out_chs)) |
|
|
|
nn.init.constant_(self.bn.weight, bn_weight_init) |
|
|
|
@torch.no_grad() |
|
def fuse(self): |
|
c, bn = self._modules.values() |
|
w = bn.weight / (bn.running_var + bn.eps) ** 0.5 |
|
w = c.weight * w[:, None, None, None] |
|
b = bn.bias - bn.running_mean * bn.weight / (bn.running_var + bn.eps) ** 0.5 |
|
m = nn.ConvTranspose2d( |
|
w.size(1), w.size(0), w.shape[2:], stride=self.c.stride, |
|
padding=self.c.padding, dilation=self.c.dilation, groups=self.c.groups) |
|
m.weight.data.copy_(w) |
|
m.bias.data.copy_(b) |
|
return m |
|
|
|
|
|
def stem_b4_transpose(in_chs, out_chs, activation): |
|
""" |
|
Modification of |
|
https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/levit.py: stem_b16 |
|
such that ConvTranspose2d is used instead of Conv2d and stem is also reduced to the half. |
|
""" |
|
return nn.Sequential( |
|
ConvTransposeNorm(in_chs, out_chs, 3, 2, 1), |
|
activation(), |
|
ConvTransposeNorm(out_chs, out_chs // 2, 3, 2, 1), |
|
activation()) |
|
|
|
|
|
def _make_pretrained_levit_384(pretrained, hooks=None): |
|
model = timm.create_model("levit_384", pretrained=pretrained) |
|
|
|
hooks = [3, 11, 21] if hooks == None else hooks |
|
return _make_levit_backbone( |
|
model, |
|
hooks=hooks |
|
) |
|
|
