wehub-resource-sync
37 lines
1.1 KiB
Python
37 lines
1.1 KiB
Python
import torch.nn as nn
|
|
from .fairseq_dropout import FairseqDropout
|
|
|
|
class Lora(nn.Module):
|
|
def __init__(self, u_dim, r, v_dim, shared_u=None, shared_v=None, dropout=0.0):
|
|
super().__init__()
|
|
self.u_dim = u_dim
|
|
self.d_dim = r
|
|
self.v_dim = v_dim
|
|
self.dropout_module = FairseqDropout(dropout, module_name=self.__class__.__name__)
|
|
|
|
if shared_u is not None:
|
|
assert shared_u.weight.size() == (r, u_dim)
|
|
self.linear_u = shared_u
|
|
else:
|
|
self.linear_u = nn.Linear(u_dim, r)
|
|
|
|
if shared_v is not None:
|
|
assert shared_v.weight.size() == (v_dim, r)
|
|
self.linear_v = shared_v
|
|
else:
|
|
self.linear_v = nn.Linear(r, v_dim)
|
|
|
|
def forward(self, x):
|
|
x = self.linear_u(x)
|
|
x = self.dropout_module(x)
|
|
x = self.linear_v(x)
|
|
x = self.dropout_module(x)
|
|
return x
|
|
|
|
@classmethod
|
|
def ratio_r(cls, layerid): # only consider 6 layers: 2 1.5 1 0.5 0.5 0.5
|
|
if layerid < 4:
|
|
return 2 - 0.5 * layerid
|
|
else:
|
|
return 0.5
|