114 lines
4.5 KiB
Python
114 lines
4.5 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright 2022-2023 The Alibaba Fundamental Vision Team Authors. All rights reserved.
|
|
from typing import Union
|
|
|
|
from torch import nn
|
|
|
|
from modelscope.metainfo import Trainers
|
|
from modelscope.models.base import Model, TorchModel
|
|
from modelscope.trainers.builder import TRAINERS
|
|
from modelscope.trainers.trainer import EpochBasedTrainer
|
|
|
|
|
|
@TRAINERS.register_module(module_name=Trainers.vision_efficient_tuning)
|
|
class VisionEfficientTuningTrainer(EpochBasedTrainer):
|
|
""" Vision Efficient Tuning Trainer based on EpochBasedTrainer
|
|
|
|
The trainer freezes the parameters of the pre-trained model and
|
|
tunes the extra parameters of the different parameter-efficient
|
|
transfer learning (PETL) method.
|
|
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def build_model(self) -> Union[nn.Module, TorchModel]:
|
|
""" Instantiate a pytorch model and return.
|
|
|
|
By default, we will create a model using config from configuration file. You can
|
|
override this method in a subclass.
|
|
|
|
"""
|
|
model = Model.from_pretrained(self.model_dir, cfg_dict=self.cfg)
|
|
if 'freeze_cfg' in self.cfg['model']:
|
|
model = self.freeze(model, **self.cfg['model']['freeze_cfg'])
|
|
if not isinstance(model, nn.Module) and hasattr(model, 'model'):
|
|
return model.model
|
|
elif isinstance(model, nn.Module):
|
|
return model
|
|
|
|
def train(self, *args, **kwargs):
|
|
self.print_model_params_status()
|
|
super().train(*args, **kwargs)
|
|
|
|
def evaluate(self, *args, **kwargs):
|
|
metric_values = super().evaluate(*args, **kwargs)
|
|
return metric_values
|
|
|
|
def freeze(self, model, freeze_part=[], train_part=[]):
|
|
""" Freeze or train the model based on the config.
|
|
|
|
Args:
|
|
model: the current model.
|
|
freeze_part: the config of frozen parameters.
|
|
train_part: the config of trainable parameters.
|
|
"""
|
|
if hasattr(model, 'module'):
|
|
freeze_model = model.module
|
|
else:
|
|
freeze_model = model
|
|
|
|
if freeze_part and len(freeze_part) > 0:
|
|
if 'backbone' in freeze_part:
|
|
part = freeze_part['backbone']
|
|
for name, param in freeze_model.model.backbone.named_parameters(
|
|
):
|
|
freeze_flag = sum([p in name for p in part]) > 0
|
|
if freeze_flag:
|
|
param.requires_grad = False
|
|
elif 'head' in freeze_part:
|
|
part = freeze_part['head']
|
|
for name, param in freeze_model.model.head.named_parameters():
|
|
freeze_flag = sum([p in name for p in part]) > 0
|
|
if freeze_flag:
|
|
param.requires_grad = False
|
|
|
|
if train_part and len(train_part) > 0:
|
|
if 'backbone' in train_part:
|
|
part = train_part['backbone']
|
|
for name, param in freeze_model.model.backbone.named_parameters(
|
|
):
|
|
freeze_flag = sum([p in name for p in part]) > 0
|
|
if freeze_flag:
|
|
param.requires_grad = True
|
|
elif 'head' in train_part:
|
|
part = train_part['head']
|
|
for name, param in freeze_model.model.head.named_parameters():
|
|
freeze_flag = sum([p in name for p in part]) > 0
|
|
if freeze_flag:
|
|
param.requires_grad = True
|
|
return model
|
|
|
|
def print_model_params_status(self, model=None, logger=None):
|
|
"""Print the status and parameters of the model"""
|
|
if model is None:
|
|
model = self.model
|
|
if logger is None:
|
|
logger = self.logger
|
|
train_param_dict = {}
|
|
all_param_numel = 0
|
|
for key, val in model.named_parameters():
|
|
if val.requires_grad:
|
|
sub_key = '.'.join(key.split('.', 1)[-1].split('.', 2)[:2])
|
|
if sub_key in train_param_dict:
|
|
train_param_dict[sub_key] += val.numel()
|
|
else:
|
|
train_param_dict[sub_key] = val.numel()
|
|
all_param_numel += val.numel()
|
|
train_param_numel = sum(train_param_dict.values())
|
|
logger.info(
|
|
f'Load trainable params {train_param_numel} / {all_param_numel} = '
|
|
f'{train_param_numel/all_param_numel:.2%}, '
|
|
f'train part: {train_param_dict}.')
|