class RichBoyNet(nn.Module):def __init__(self, hidden_size_1=1000, hidden_size_2=2000):super(RichBoyNet,self).__init__()self.linear1 = nn.Linear(28*28, hidden_size_1) self.linear2 = nn.Linear(hidden_size_1, hidden_size_2) self.linear3 = nn.Linear(hidden_size_2, 10)self.relu = nn.ReLU()
def forward(self, img):x = img.view(-1, 28*28)x = self.relu(self.linear1(x))x = self.relu(self.linear2(x))x = self.linear3(x)return x
net = RichBoyNet().to(device)

class LoRAParametrization(nn.Module):def __init__(self, features_in, features_out, rank=1, alpha=1, device='cpu'):super().__init__()# Section 4.1 of the paper: # We use a random Gaussian initialization for A and zero for B, so ∆W = BA is zero at the beginning of trainingself.lora_A = nn.Parameter(torch.zeros((rank,features_out)).to(device))self.lora_B = nn.Parameter(torch.zeros((features_in, rank)).to(device))nn.init.normal_(self.lora_A, mean=0, std=1)# Section 4.1 of the paper: # We then scale ∆Wx by α/r , where α is a constant in r. # When optimizing with Adam, tuning α is roughly the same as tuning the learning rate if we scale the initialization appropriately. # As a result, we simply set α to the first r we try and do not tune it. # This scaling helps to reduce the need to retune hyperparameters when we vary r.self.scale = alpha / rankself.enabled = True
def forward(self, original_weights):if self.enabled:# Return W + (B*A)*scalereturn original_weights + torch.matmul(self.lora_B, self.lora_A).view(original_weights.shape) * self.scaleelse:return original_weights