wehub-resource-sync
204 lines
9.1 KiB
Python
204 lines
9.1 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
|
|
import re
|
|
import torch
|
|
import types
|
|
from dataclasses import dataclass, field
|
|
from torch import nn
|
|
from typing import List, Union
|
|
|
|
from swift.utils import find_sub_module, get_logger
|
|
from .utils import ActivationMixin, SwiftAdapter, SwiftConfig, SwiftOutput
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
@dataclass
|
|
class PromptConfig(SwiftConfig):
|
|
"""
|
|
The configuration class for the prompt module.
|
|
|
|
Visual prompt tuning (VPT) is proposed to initialize tunable prompt tokens
|
|
and prepend to the original tokens in the first layer or multiple layers.
|
|
'Visual Prompt Tuning' by Jia et al.(2022)
|
|
See https://arxiv.org/abs/2203.12119
|
|
|
|
Here we apply the VPT to other fields.
|
|
|
|
Args:
|
|
dim(`Union[int, List[int]]`): The dimension of the hidden states, use list if there are up-sample blocks
|
|
or down-sample blocks
|
|
target_modules(str): The layer module to be replaced, in regex format
|
|
embedding_pos(Union[str, int]): The position of the embedding tensor
|
|
attention_mask_pos(Union[str, int]): The position of the attention mask
|
|
attention_mask_value(Union[float, int, bool]): The value to pad to the attention mask
|
|
prompt_length(int): The length of the prompt tokens
|
|
attach_front(bool): When set to True, prompt is attached in front of the embedding
|
|
extract_embedding(bool): Whether the embedding is extracted at final stage to keep the same dims with inputs
|
|
"""
|
|
|
|
dim: Union[int, List[int]] = field(default=None, metadata={'help': 'The dimension of the hidden states'})
|
|
|
|
target_modules: str = field(default=None, metadata={'help': 'The layer module to be replaced, in regex format'})
|
|
|
|
embedding_pos: Union[str, int] = field(default=None, metadata={'help': 'The position of the embedding tensor'})
|
|
|
|
attention_mask_pos: Union[str, int] = field(default=None, metadata={'help': 'The position of the attention mask'})
|
|
|
|
attention_mask_value: Union[float, int, bool] = field(
|
|
default=0., metadata={'help': 'The value to pad to the attention mask'})
|
|
|
|
prompt_length: int = field(default=16, metadata={'help': 'The length of the prompt tokens'})
|
|
|
|
attach_front: bool = field(
|
|
default=True, metadata={'help': 'When set to True, prompt is attached in front of the embedding'})
|
|
|
|
extract_embedding: bool = field(
|
|
default=False,
|
|
metadata={'help': 'Whether the embedding is extracted at final stage to keep the same dims with inputs'})
|
|
|
|
def __post_init__(self):
|
|
from .mapping import SwiftTuners
|
|
self.swift_type = SwiftTuners.PROMPT
|
|
|
|
|
|
class Prompt(SwiftAdapter):
|
|
|
|
@staticmethod
|
|
def prepare_model(model: nn.Module, config: PromptConfig, adapter_name: str):
|
|
module_keys = [key for key, _ in model.named_modules()]
|
|
match_module_keys = []
|
|
for module_key in module_keys:
|
|
if isinstance(config.target_modules, str):
|
|
target_module_found = re.fullmatch(config.target_modules, module_key)
|
|
else:
|
|
target_module_found = any(module_key.endswith(target_key) for target_key in config.target_modules)
|
|
if target_module_found: # noqa
|
|
module = model.get_submodule(module_key)
|
|
|
|
def _forward(self, *args, **kwargs):
|
|
if isinstance(config.embedding_pos, int):
|
|
input_embedding = args[config.embedding_pos]
|
|
else:
|
|
input_embedding = kwargs[config.embedding_pos]
|
|
|
|
input_embedding = getattr(self, f'prompt_{adapter_name}').forward(input_embedding)
|
|
if isinstance(config.embedding_pos, int):
|
|
args = type(args)(
|
|
args[0:config.embedding_pos] + (input_embedding, ) + args[config.embedding_pos + 1:])
|
|
else:
|
|
kwargs[config.embedding_pos] = input_embedding
|
|
|
|
if config.attention_mask_pos:
|
|
attention_mask = None
|
|
if isinstance(config.attention_mask_pos, int):
|
|
attention_mask = args[config.attention_mask_pos]
|
|
elif isinstance(config.attention_mask_pos, str):
|
|
attention_mask = kwargs[config.attention_mask_pos]
|
|
|
|
if attention_mask is not None:
|
|
attention_mask = getattr(self,
|
|
f'prompt_{adapter_name}').patch_attention_mask(attention_mask)
|
|
if isinstance(config.attention_mask_pos, int):
|
|
args = type(args)(
|
|
args[0:config.attention_mask_pos] + (attention_mask, )
|
|
+ args[config.attention_mask_pos + 1:])
|
|
else:
|
|
kwargs[config.attention_mask_pos] = attention_mask
|
|
|
|
forward_output = getattr(self, f'forward_origin_{adapter_name}')(*args, **kwargs)
|
|
if config.extract_embedding:
|
|
forward_output = getattr(self, f'prompt_{adapter_name}').extract(forward_output)
|
|
|
|
return forward_output
|
|
|
|
setattr(module, f'forward_origin_{adapter_name}', module.forward)
|
|
module.forward = types.MethodType(_forward, module)
|
|
if isinstance(config.dim, list):
|
|
input_dim = config.dim[len(match_module_keys)]
|
|
else:
|
|
input_dim = config.dim
|
|
prompt_module = PromptModule(input_dim, int(module_key.rsplit('.')[-1]), adapter_name, module_key,
|
|
config.prompt_length, config.attention_mask_value, config.attach_front)
|
|
setattr(module, f'prompt_{adapter_name}', prompt_module)
|
|
logger.info(f'Prompt modules(module_key): {module_key}.prompt_{adapter_name}')
|
|
match_module_keys.append(module_key)
|
|
|
|
def state_dict_callback(state_dict, adapter_name, **kwargs):
|
|
return {key: value for key, value in state_dict.items() if f'prompt_{adapter_name}' in key}
|
|
|
|
def mark_trainable_callback(model):
|
|
return
|
|
|
|
return SwiftOutput(
|
|
config=config, state_dict_callback=state_dict_callback, mark_trainable_callback=mark_trainable_callback)
|
|
|
|
@staticmethod
|
|
def activate_adapter(module: torch.nn.Module, adapter_name: str, activate: bool, offload: str = None):
|
|
modules = find_sub_module(module, f'prompt_{adapter_name}')
|
|
for _module in modules:
|
|
_module: ActivationMixin
|
|
_module: nn.Module
|
|
_module.set_activation(adapter_name, activate)
|
|
SwiftAdapter.save_memory(_module, adapter_name, _module.module_key, activate, offload)
|
|
|
|
|
|
class PromptModule(nn.Module, ActivationMixin):
|
|
"""The implementation of vision prompt tuning method.
|
|
|
|
Visual prompt tuning (VPT) is proposed to initialize tunable prompt tokens
|
|
and prepend to the original tokens in the first layer or multiple layers.
|
|
'Visual Prompt Tuning' by Jia et al.(2022)
|
|
See https://arxiv.org/abs/2203.12119
|
|
|
|
Args:
|
|
dim: An integer indicating the embedding dimension.
|
|
layer_num: An integer indicating number of layers.
|
|
prompt_length: An integer indicating the length of vision prompt tuning.
|
|
"""
|
|
|
|
def __init__(self, dim, layer_num, adapter_name, module_key, prompt_length=None, mask_values=0., attach_front=True):
|
|
super(PromptModule, self).__init__()
|
|
super(nn.Module, self).__init__(module_key)
|
|
self.dim = dim
|
|
self.layer_num = layer_num
|
|
self.adapter_name = adapter_name
|
|
self.prompt_length = prompt_length
|
|
self.mask_values = mask_values
|
|
self.attach_front = attach_front
|
|
self.prompt_token = nn.Parameter(torch.zeros(1, prompt_length, dim))
|
|
nn.init.xavier_uniform_(self.prompt_token)
|
|
self.mark_all_sub_modules_as_plugin()
|
|
|
|
def forward(self, x):
|
|
if not self.is_activated(self.adapter_name):
|
|
return x
|
|
prompt_token = self.prompt_token.expand(x.shape[0], -1, -1).to(x.device, x.dtype)
|
|
|
|
if self.layer_num == 0:
|
|
if self.attach_front:
|
|
x = torch.cat((prompt_token, x), dim=1)
|
|
else:
|
|
x = torch.cat((x, prompt_token), dim=1)
|
|
else:
|
|
if self.attach_front:
|
|
x = torch.cat((prompt_token, x[:, self.prompt_length:, :]), dim=1)
|
|
else:
|
|
x = torch.cat((x[:, :-self.prompt_length, :], prompt_token), dim=1)
|
|
return x
|
|
|
|
def patch_attention_mask(self, m):
|
|
if not self.is_activated(self.adapter_name):
|
|
return m
|
|
prefix_attention_mask = torch.full((*m.shape[:-1], self.prompt_length), self.mask_values).to(m.device)
|
|
if self.attach_front:
|
|
return torch.cat((prefix_attention_mask, m), dim=-1)
|
|
else:
|
|
return torch.cat((m, prefix_attention_mask), dim=-1)
|
|
|
|
def extract(self, x):
|
|
if self.attach_front:
|
|
return x[:, self.prompt_length:, :]
|
|
else:
|
|
return x[:, :-self.prompt_length, :]
|