File size: 3,922 Bytes
7feac49 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
# Copyright (c) Alibaba, Inc. and its affiliates.
import math
import torch
import torch.nn as nn
from swift.utils.logger import get_logger
logger = get_logger()
def detach_tensors(feats):
if type(feats) in [list, tuple]:
feats = [detach_tensors(feat) if feat is not None else None for feat in feats]
elif isinstance(feats, dict):
feats = {key: detach_tensors(val) for key, val in feats.items()}
elif isinstance(feats, torch.Tensor):
feats = feats.detach()
else:
feats = feats.detach()
return feats
def probe_tensors(module, feats, name):
feats = detach_tensors(feats)
setattr(module, name, feats)
def probe_input_pre_hook(self, args):
input = args[0]
probe_tensors(self, input, 'probe_input_data')
return args
def probe_output_hook(self, args, result):
output = result
probe_tensors(self, output, 'probe_output_data')
return output
def choose_weight_type(weight_type, dim):
if weight_type == 'gate':
scaling = nn.Linear(dim, 1)
elif weight_type == 'scale':
scaling = nn.Parameter(torch.Tensor(1))
scaling.data.fill_(1)
elif weight_type == 'scale_channel':
scaling = nn.Parameter(torch.Tensor(dim))
scaling.data.fill_(1)
elif weight_type and weight_type.startswith('scalar'):
scaling = float(weight_type.split('_')[-1])
else:
scaling = None
return scaling
def get_weight_value(weight_type, scaling, x):
if weight_type in ['gate']:
scaling = torch.mean(torch.sigmoid(scaling(x)), dim=1).view(-1, 1, 1)
elif weight_type in ['scale', 'scale_channel'] or weight_type.startswith('scalar'):
scaling = scaling
else:
scaling = None
return scaling
class SCEAdapter(nn.Module):
def __init__(self,
dim,
adapter_length,
adapter_type=None,
adapter_weight=None,
act_layer=nn.GELU,
zero_init_last=True,
use_bias=True):
super(SCEAdapter, self).__init__()
self.dim = dim
self.adapter_length = adapter_length
self.adapter_type = adapter_type
self.adapter_weight = adapter_weight
self.zero_init_last = zero_init_last
self.ln1 = nn.Linear(dim, adapter_length, bias=use_bias)
self.activate = act_layer()
self.ln2 = nn.Linear(adapter_length, dim, bias=use_bias)
self.init_weights()
self.init_scaling()
def _zero_init_weights(self, m):
if isinstance(m, nn.Linear):
nn.init.zeros_(m.weight)
nn.init.zeros_(m.bias)
def _kaiming_init_weights(self, m):
if isinstance(m, nn.Linear):
nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5))
def init_weights(self):
self._kaiming_init_weights(self.ln1)
if self.zero_init_last:
self._zero_init_weights(self.ln2)
else:
self._kaiming_init_weights(self.ln2)
def init_scaling(self):
if self.adapter_weight:
self.scaling = choose_weight_type(self.adapter_weight, self.dim)
else:
self.scaling = None
def forward(self, x, x_shortcut=None, use_shortcut=True, **kwargs):
if x_shortcut is None:
x_shortcut = x
x_shape = x.shape
if len(x_shape) == 4:
b, d, h, w = x_shape
x = x.permute(0, 2, 3, 1).reshape(b, h * w, d)
out = self.ln2(self.activate(self.ln1(x)))
if self.adapter_weight:
scaling = get_weight_value(self.adapter_weight, self.scaling, out)
out = out * scaling if scaling is not None else out
if len(x_shape) == 4:
b, d, h, w = x_shape
out = out.reshape(b, h, w, -1).permute(0, 3, 1, 2).contiguous()
if use_shortcut:
out = x_shortcut + out
return out
|