73 lines
2.4 KiB
Python
73 lines
2.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import math
|
|
|
|
import torch.nn as nn
|
|
|
|
from scepter.modules.model.tuner.tuner_utils import (choose_weight_type,
|
|
get_weight_value)
|
|
|
|
|
|
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.use_bias = use_bias
|
|
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)
|
|
if self.use_bias:
|
|
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
|