LoRA实现

使用PyTorch实现LoRA并加载到网络中的示例代码

使用setattr和getattr来动态地修改模型，插入LoRA层。
定义LoRA层：

import torch
import torch.nn as nn
class LoRALayer(nn.Module):
    def __init__(self, in_features, out_features, rank, scaling=1.0):
        super(LoRALayer, self).__init__()
        self.rank = rank
        self.scaling = scaling
        self.A = nn.Parameter(torch.randn(in_features, rank))
        self.B = nn.Parameter(torch.randn(rank, out_features))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.normal_(self.A, std=0.02)
        nn.init.zeros_(self.B)

    def forward(self, x):
        return x + self.scaling * torch.matmul(torch.matmul(x, self.A), self.B)

定义一个函数将LoRA层插入到预训练模型中：

def insert_lora(model, rank, scaling):
    model_dict = [name for name in model.named_modules()]
    # print(model_dict)
    for name, module in model_dict:
        print("name, module",name, module)
        if isinstance(module, nn.Linear):  # 只针对全连接层插入LoRA
            lora_layer = LoRALayer(module.in_features, module.out_features, rank, scaling)
            setattr(model, f'{name}_lora', lora_layer)
            # 重写前向传播以包含LoRA层
            original_forward = module.forward
            def new_forward(x):
                x = original_forward(x)
                lora = getattr(model, f'{name}_lora')
                return lora(x)
            module.forward = new_forward.__get__(module, type(module))

创建一个简单的网络，应用LoRA：

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(10, 20)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(20, 5)
    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x
# 创建模型
model = SimpleModel()
print("simple model: ",model)
# 应用LoRA
insert_lora(model, rank=4, scaling=1.0)
# 现在模型包含了LoRA层
print("LoRa model: ",model)

最后，训练模型，只需要更新LoRA层的参数：

# 假设我们有一个数据加载器
# dataloader = ...
# 定义优化器，只优化LoRA层的参数
lora_parameters = [p for n, p in model.named_parameters() if 'lora' in n]
optimizer = torch.optim.Adam(lora_parameters, lr=3e-4)
# 训练代码
for epoch in range(num_epochs):
    for inputs, targets in dataloader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = nn.functional.mse_loss(outputs, targets)
        loss.backward()
        optimizer.step()