| import torch |
|
|
| import torch.nn as nn |
|
|
|
|
| class Slice(nn.Module): |
| def __init__(self, start_index=1): |
| super(Slice, self).__init__() |
| self.start_index = start_index |
|
|
| def forward(self, x): |
| return x[:, self.start_index:] |
|
|
|
|
| class AddReadout(nn.Module): |
| def __init__(self, start_index=1): |
| super(AddReadout, self).__init__() |
| self.start_index = start_index |
|
|
| def forward(self, x): |
| if self.start_index == 2: |
| readout = (x[:, 0] + x[:, 1]) / 2 |
| else: |
| readout = x[:, 0] |
| return x[:, self.start_index:] + readout.unsqueeze(1) |
|
|
|
|
| class ProjectReadout(nn.Module): |
| def __init__(self, in_features, start_index=1): |
| super(ProjectReadout, self).__init__() |
| self.start_index = start_index |
|
|
| self.project = nn.Sequential(nn.Linear(2 * in_features, in_features), nn.GELU()) |
|
|
| def forward(self, x): |
| readout = x[:, 0].unsqueeze(1).expand_as(x[:, self.start_index:]) |
| features = torch.cat((x[:, self.start_index:], readout), -1) |
|
|
| return self.project(features) |
|
|
|
|
| class Transpose(nn.Module): |
| def __init__(self, dim0, dim1): |
| super(Transpose, self).__init__() |
| self.dim0 = dim0 |
| self.dim1 = dim1 |
|
|
| def forward(self, x): |
| x = x.transpose(self.dim0, self.dim1) |
| return x |
|
|
|
|
| activations = {} |
|
|
|
|
| def get_activation(name): |
| def hook(model, input, output): |
| activations[name] = output |
|
|
| return hook |
|
|
|
|
| def forward_default(pretrained, x, function_name="forward_features"): |
| exec(f"pretrained.model.{function_name}(x)") |
|
|
| layer_1 = pretrained.activations["1"] |
| layer_2 = pretrained.activations["2"] |
| layer_3 = pretrained.activations["3"] |
| layer_4 = pretrained.activations["4"] |
|
|
| if hasattr(pretrained, "act_postprocess1"): |
| layer_1 = pretrained.act_postprocess1(layer_1) |
| if hasattr(pretrained, "act_postprocess2"): |
| layer_2 = pretrained.act_postprocess2(layer_2) |
| if hasattr(pretrained, "act_postprocess3"): |
| layer_3 = pretrained.act_postprocess3(layer_3) |
| if hasattr(pretrained, "act_postprocess4"): |
| layer_4 = pretrained.act_postprocess4(layer_4) |
|
|
| return layer_1, layer_2, layer_3, layer_4 |
|
|
|
|
| def forward_adapted_unflatten(pretrained, x, function_name="forward_features"): |
| b, c, h, w = x.shape |
|
|
| exec(f"glob = pretrained.model.{function_name}(x)") |
|
|
| layer_1 = pretrained.activations["1"] |
| layer_2 = pretrained.activations["2"] |
| layer_3 = pretrained.activations["3"] |
| layer_4 = pretrained.activations["4"] |
|
|
| layer_1 = pretrained.act_postprocess1[0:2](layer_1) |
| layer_2 = pretrained.act_postprocess2[0:2](layer_2) |
| layer_3 = pretrained.act_postprocess3[0:2](layer_3) |
| layer_4 = pretrained.act_postprocess4[0:2](layer_4) |
|
|
| unflatten = nn.Sequential( |
| nn.Unflatten( |
| 2, |
| torch.Size( |
| [ |
| h // pretrained.model.patch_size[1], |
| w // pretrained.model.patch_size[0], |
| ] |
| ), |
| ) |
| ) |
|
|
| if layer_1.ndim == 3: |
| layer_1 = unflatten(layer_1) |
| if layer_2.ndim == 3: |
| layer_2 = unflatten(layer_2) |
| if layer_3.ndim == 3: |
| layer_3 = unflatten(layer_3) |
| if layer_4.ndim == 3: |
| layer_4 = unflatten(layer_4) |
|
|
| layer_1 = pretrained.act_postprocess1[3: len(pretrained.act_postprocess1)](layer_1) |
| layer_2 = pretrained.act_postprocess2[3: len(pretrained.act_postprocess2)](layer_2) |
| layer_3 = pretrained.act_postprocess3[3: len(pretrained.act_postprocess3)](layer_3) |
| layer_4 = pretrained.act_postprocess4[3: len(pretrained.act_postprocess4)](layer_4) |
|
|
| return layer_1, layer_2, layer_3, layer_4 |
|
|
|
|
| def get_readout_oper(vit_features, features, use_readout, start_index=1): |
| if use_readout == "ignore": |
| readout_oper = [Slice(start_index)] * len(features) |
| elif use_readout == "add": |
| readout_oper = [AddReadout(start_index)] * len(features) |
| elif use_readout == "project": |
| readout_oper = [ |
| ProjectReadout(vit_features, start_index) for out_feat in features |
| ] |
| else: |
| assert ( |
| False |
| ), "wrong operation for readout token, use_readout can be 'ignore', 'add', or 'project'" |
|
|
| return readout_oper |
|
|
|
|
| def make_backbone_default( |
| model, |
| features=[96, 192, 384, 768], |
| size=[384, 384], |
| hooks=[2, 5, 8, 11], |
| vit_features=768, |
| use_readout="ignore", |
| start_index=1, |
| start_index_readout=1, |
| ): |
| 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.model.blocks[hooks[3]].register_forward_hook(get_activation("4")) |
|
|
| pretrained.activations = activations |
|
|
| readout_oper = get_readout_oper(vit_features, features, use_readout, start_index_readout) |
|
|
| |
| pretrained.act_postprocess1 = nn.Sequential( |
| readout_oper[0], |
| Transpose(1, 2), |
| nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), |
| nn.Conv2d( |
| in_channels=vit_features, |
| out_channels=features[0], |
| kernel_size=1, |
| stride=1, |
| padding=0, |
| ), |
| nn.ConvTranspose2d( |
| in_channels=features[0], |
| out_channels=features[0], |
| kernel_size=4, |
| stride=4, |
| padding=0, |
| bias=True, |
| dilation=1, |
| groups=1, |
| ), |
| ) |
|
|
| pretrained.act_postprocess2 = nn.Sequential( |
| readout_oper[1], |
| Transpose(1, 2), |
| nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), |
| nn.Conv2d( |
| in_channels=vit_features, |
| out_channels=features[1], |
| kernel_size=1, |
| stride=1, |
| padding=0, |
| ), |
| nn.ConvTranspose2d( |
| in_channels=features[1], |
| out_channels=features[1], |
| kernel_size=2, |
| stride=2, |
| padding=0, |
| bias=True, |
| dilation=1, |
| groups=1, |
| ), |
| ) |
|
|
| pretrained.act_postprocess3 = nn.Sequential( |
| readout_oper[2], |
| Transpose(1, 2), |
| nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), |
| nn.Conv2d( |
| in_channels=vit_features, |
| out_channels=features[2], |
| kernel_size=1, |
| stride=1, |
| padding=0, |
| ), |
| ) |
|
|
| pretrained.act_postprocess4 = nn.Sequential( |
| readout_oper[3], |
| Transpose(1, 2), |
| nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), |
| nn.Conv2d( |
| in_channels=vit_features, |
| out_channels=features[3], |
| kernel_size=1, |
| stride=1, |
| padding=0, |
| ), |
| nn.Conv2d( |
| in_channels=features[3], |
| out_channels=features[3], |
| kernel_size=3, |
| stride=2, |
| padding=1, |
| ), |
| ) |
|
|
| pretrained.model.start_index = start_index |
| pretrained.model.patch_size = [16, 16] |
|
|
| return pretrained |
|
|
