95 lines
3.3 KiB
Python
95 lines
3.3 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright 2021-2022 The Alibaba Fundamental Vision Team Authors. All rights reserved.
|
|
|
|
import os
|
|
from typing import Any, Dict
|
|
|
|
import json
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from modelscope.metainfo import Models
|
|
from modelscope.models.base.base_torch_model import TorchModel
|
|
from modelscope.models.builder import MODELS
|
|
from modelscope.outputs import OutputKeys
|
|
from modelscope.utils.constant import ModelFile, Tasks
|
|
from .backbone import CLIP, ProbingModel
|
|
|
|
|
|
@MODELS.register_module(
|
|
Tasks.image_classification, module_name=Models.image_probing_model)
|
|
class StructuredProbingModel(TorchModel):
|
|
"""
|
|
The implementation of 'Structured Model Probing: Empowering
|
|
Efficient Adaptation by Structured Regularization'.
|
|
"""
|
|
|
|
def __init__(self, model_dir, *args, **kwargs):
|
|
"""
|
|
Initialize a probing model.
|
|
Args:
|
|
model_dir: model id or path
|
|
"""
|
|
super(StructuredProbingModel, self).__init__()
|
|
model_dir = os.path.join(model_dir, 'food101-clip-vitl14-full.pt')
|
|
model_file = torch.load(model_dir)
|
|
self.feature_size = model_file['meta_info']['feature_size']
|
|
self.num_classes = model_file['meta_info']['num_classes']
|
|
self.backbone = CLIP(
|
|
'CLIP_ViTL14_FP16',
|
|
use_pretrain=True,
|
|
state_dict=model_file['backbone_model_state_dict'])
|
|
self.probing_model = ProbingModel(self.feature_size, self.num_classes)
|
|
self.probing_model.load_state_dict(
|
|
model_file['probing_model_state_dict'])
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Forward Function of SMP.
|
|
Args:
|
|
x: the input images (B, 3, H, W)
|
|
"""
|
|
|
|
keys = []
|
|
for idx in range(0, 24):
|
|
keys.append('layer_{}_pre_attn'.format(idx))
|
|
keys.append('layer_{}_attn'.format(idx))
|
|
keys.append('layer_{}_mlp'.format(idx))
|
|
keys.append('pre_logits')
|
|
features = self.backbone(x.half())
|
|
features_agg = []
|
|
for i in keys:
|
|
aggregated_feature = self.aggregate_token(features[i], 1024)
|
|
features_agg.append(aggregated_feature)
|
|
features_agg = torch.cat((features_agg), dim=1)
|
|
outputs = self.probing_model(features_agg.float())
|
|
return outputs
|
|
|
|
def aggregate_token(self, output, target_size):
|
|
"""
|
|
Aggregating features from tokens.
|
|
Args:
|
|
output: the output of intermidiant features
|
|
from a ViT model
|
|
target_size: target aggregated feature size
|
|
"""
|
|
if len(output.shape) == 3:
|
|
_, n_token, channels = output.shape
|
|
if channels >= target_size:
|
|
pool_size = 0
|
|
else:
|
|
n_groups = target_size / channels
|
|
pool_size = int(n_token / n_groups)
|
|
|
|
if pool_size > 0:
|
|
output = torch.permute(output, (0, 2, 1))
|
|
output = torch.nn.AvgPool1d(
|
|
kernel_size=pool_size, stride=pool_size)(
|
|
output)
|
|
output = torch.flatten(output, start_dim=1)
|
|
else:
|
|
output = torch.mean(output, dim=1)
|
|
output = torch.nn.functional.normalize(output, dim=1)
|
|
return output
|