276 lines
9.2 KiB
Python
276 lines
9.2 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright 2022-2023 The Alibaba Fundamental Vision Team Authors. All rights reserved.
|
|
import math
|
|
from collections import OrderedDict
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torchvision
|
|
|
|
|
|
class Prompt(nn.Module):
|
|
"""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
|
|
|
|
Attributes:
|
|
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.
|
|
prompt_type: A string indicating the type of vision prompt tuning.
|
|
"""
|
|
|
|
def __init__(self, dim, layer_num, prompt_length=None, prompt_type=None):
|
|
super(Prompt, self).__init__()
|
|
self.dim = dim
|
|
self.layer_num = layer_num
|
|
self.prompt_length = prompt_length
|
|
self.prompt_type = prompt_type
|
|
|
|
self.prompt_token = nn.Parameter(torch.zeros(1, prompt_length, dim))
|
|
nn.init.xavier_uniform_(self.prompt_token)
|
|
|
|
def forward(self, x):
|
|
B, N, C = x.shape
|
|
prompt_token = self.prompt_token.expand(B, -1, -1)
|
|
|
|
if self.layer_num == 0:
|
|
x = torch.cat((x, prompt_token), dim=1)
|
|
else:
|
|
x = torch.cat((x[:, :-self.prompt_length, :], prompt_token), dim=1)
|
|
return x
|
|
|
|
|
|
class Adapter(nn.Module):
|
|
"""The implementation of adapter tuning method.
|
|
|
|
Adapters project input tokens by an MLP layer.
|
|
'Parameter-Efficient Transfer Learning for NLP' by Houlsby et al.(2019)
|
|
See http://arxiv.org/abs/1902.00751
|
|
|
|
Attributes:
|
|
dim: An integer indicating the embedding dimension.
|
|
adapter_length: An integer indicating the length of adapter tuning.
|
|
adapter_type: A string indicating the type of adapter tuning.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim,
|
|
adapter_length=None,
|
|
adapter_type=None,
|
|
act_layer=nn.GELU,
|
|
):
|
|
super(Adapter, self).__init__()
|
|
self.dim = dim
|
|
self.adapter_length = adapter_length
|
|
self.adapter_type = adapter_type
|
|
self.ln1 = nn.Linear(dim, adapter_length)
|
|
self.activate = act_layer()
|
|
self.ln2 = nn.Linear(adapter_length, dim)
|
|
self.init_weights()
|
|
|
|
def init_weights(self):
|
|
|
|
def _init_weights(m):
|
|
if isinstance(m, nn.Linear):
|
|
nn.init.xavier_uniform_(m.weight)
|
|
nn.init.normal_(m.bias, std=1e-6)
|
|
|
|
self.apply(_init_weights)
|
|
|
|
def forward(self, x, identity=None):
|
|
out = self.ln2(self.activate(self.ln1(x)))
|
|
if identity is None:
|
|
identity = x
|
|
out = identity + out
|
|
return out
|
|
|
|
|
|
class LoRA(nn.Module):
|
|
"""The implementation of LoRA tuning method.
|
|
|
|
LoRA constructs an additional layer with low-rank decomposition matrices of the weights in the network.
|
|
'LoRA: Low-Rank Adaptation of Large Language Models' by Hu et al.(2021)
|
|
See https://arxiv.org/abs/2106.09685
|
|
|
|
Attributes:
|
|
dim: An integer indicating the embedding dimension.
|
|
num_heads: An integer indicating number of attention heads.
|
|
lora_length: An integer indicating the length of LoRA tuning.
|
|
lora_type: A string indicating the type of LoRA tuning.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim,
|
|
num_heads,
|
|
lora_length=None,
|
|
lora_type=None,
|
|
):
|
|
super(LoRA, self).__init__()
|
|
self.dim = dim
|
|
self.num_heads = num_heads
|
|
self.lora_a = nn.Linear(dim, lora_length, bias=False)
|
|
nn.init.kaiming_uniform_(self.lora_a.weight, a=math.sqrt(5))
|
|
self.lora_b = nn.Linear(lora_length, dim * 3, bias=False)
|
|
nn.init.zeros_(self.lora_b.weight)
|
|
|
|
self.lora_length = lora_length
|
|
self.lora_type = lora_type
|
|
|
|
def forward(self, x, q, k, v):
|
|
B, N, C = x.shape
|
|
qkv_delta = self.lora_b(self.lora_a(x))
|
|
qkv_delta = qkv_delta.reshape(B, N, 3, self.num_heads,
|
|
C // self.num_heads).permute(
|
|
2, 0, 3, 1, 4)
|
|
q_delta, k_delta, v_delta = qkv_delta.unbind(0)
|
|
q, k, v = q + q_delta, k + k_delta, v + v_delta
|
|
return q, k, v
|
|
|
|
|
|
class Prefix(nn.Module):
|
|
"""The implementation of prefix tuning method.
|
|
|
|
Prefix tuning optimizes the task-specific vector in the multi-head attention layer.
|
|
'Prefix-tuning: Optimizing continuous prompts for generation' by Li & Liang(2021)
|
|
See https://arxiv.org/abs/2101.00190
|
|
|
|
Attributes:
|
|
dim: An integer indicating the embedding dimension.
|
|
num_heads: An integer indicating number of attention heads.
|
|
prefix_length: An integer indicating the length of prefix tuning.
|
|
prefix_type: A string indicating the type of prefix tuning.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim,
|
|
num_heads,
|
|
prefix_length=None,
|
|
prefix_type=None,
|
|
):
|
|
super(Prefix, self).__init__()
|
|
self.dim = dim
|
|
self.num_heads = num_heads
|
|
self.prefix_length = prefix_length
|
|
self.prefix_type = prefix_type
|
|
self.prefix_key = nn.Parameter(torch.zeros(1, prefix_length, dim))
|
|
self.prefix_value = nn.Parameter(torch.zeros(1, prefix_length, dim))
|
|
nn.init.xavier_uniform_(self.prefix_key)
|
|
nn.init.xavier_uniform_(self.prefix_value)
|
|
|
|
def forward(self, x, q, k, v):
|
|
B, N, C = x.shape
|
|
prefix_key = self.prefix_key.expand(B, -1, -1).reshape(
|
|
B, self.prefix_length, self.num_heads,
|
|
self.dim // self.num_heads).permute(0, 2, 1, 3)
|
|
prefix_value = self.prefix_value.expand(B, -1, -1).reshape(
|
|
B, self.prefix_length, self.num_heads,
|
|
self.dim // self.num_heads).permute(0, 2, 1, 3)
|
|
k, v = torch.cat((k, prefix_key), dim=2), torch.cat((v, prefix_value),
|
|
dim=2)
|
|
return q, k, v
|
|
|
|
|
|
class SideTune(nn.Module):
|
|
"""The implementation of vision side-tuning method.
|
|
|
|
Side-Tuning only needs to train one side network and
|
|
weights the output of pre-trained model and side network.
|
|
'Side-Tuning: A Baseline for Network Adaptation via Additive Side Networks'
|
|
by Zhang et al.(2019)
|
|
See https://arxiv.org/abs/1912.13503
|
|
|
|
Attributes:
|
|
sidetune_length: An integer indicating the linear dimension.
|
|
sidetune_type: A string indicating the type of side network.
|
|
"""
|
|
|
|
def __init__(self, sidetune_length=None, sidetune_type=None):
|
|
super(SideTune, self).__init__()
|
|
self.sidetune_length = sidetune_length
|
|
self.sidetune_type = sidetune_type
|
|
if sidetune_type.lower() == 'fcn4':
|
|
self.side = FCN4(out_dims=self.sidetune_length)
|
|
if sidetune_type.lower() == 'alexnet':
|
|
mm = torchvision.models.alexnet(pretrained=True)
|
|
self.side = nn.Sequential(
|
|
OrderedDict([
|
|
('features', mm.features), ('avgpool', mm.avgpool),
|
|
('flatten', nn.Flatten()),
|
|
('fc', nn.Linear(9216, self.sidetune_length, bias=False))
|
|
]))
|
|
self.alpha = nn.Parameter(torch.tensor(0.0))
|
|
|
|
def forward(self, x, x_base):
|
|
alpha_squashed = torch.sigmoid(self.alpha)
|
|
x_side = self.side(x)
|
|
x_out = alpha_squashed * x_base + (1 - alpha_squashed) * x_side
|
|
return x_out
|
|
|
|
|
|
class FCN4(nn.Module):
|
|
"""The implementation of simple FCN4 network for side network.
|
|
"""
|
|
|
|
def __init__(self, out_dims=-1, **kwargs):
|
|
super(FCN4, self).__init__(**kwargs)
|
|
|
|
self.conv1 = nn.Sequential(
|
|
nn.Conv2d(
|
|
3,
|
|
16,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1,
|
|
bias=False,
|
|
dilation=1), nn.GroupNorm(2, 16), nn.ReLU())
|
|
self.conv2 = nn.Sequential(
|
|
nn.Conv2d(
|
|
16,
|
|
16,
|
|
kernel_size=3,
|
|
stride=2,
|
|
padding=0,
|
|
bias=False,
|
|
dilation=1), nn.GroupNorm(2, 16), nn.ReLU())
|
|
self.conv3 = nn.Sequential(
|
|
nn.Conv2d(
|
|
16,
|
|
32,
|
|
kernel_size=3,
|
|
stride=2,
|
|
padding=0,
|
|
bias=False,
|
|
dilation=1), nn.GroupNorm(2, 32), nn.ReLU())
|
|
self.conv4 = nn.Sequential(
|
|
nn.Conv2d(
|
|
32,
|
|
64,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=0,
|
|
bias=False,
|
|
dilation=1), nn.GroupNorm(2, 64), nn.ReLU())
|
|
self.pool = nn.AdaptiveAvgPool2d((1, 1))
|
|
if out_dims > 0:
|
|
self.fc = nn.Linear(64, out_dims)
|
|
else:
|
|
self.fc = None
|
|
|
|
def forward(self, x):
|
|
x = self.conv1(x)
|
|
x = self.conv2(x)
|
|
x = self.conv3(x)
|
|
x = self.conv4(x)
|
|
x = self.pool(x)
|
|
x = x.view(x.size(0), -1)
|
|
if self.fc is not None:
|
|
x = self.fc(x)
|
|
return x
|