LoRA原理与实现--PyTorch自己搭建LoRA模型-灵析社区

一、前言

在AIGC领域频繁出现着一个特殊名词“LoRA”，听上去有点像人名，但是这是一种模型训练的方法。LoRA全称Low-Rank Adaptation of Large Language Models，中文叫做大语言模型的低阶适应。如今在stable diffusion中用地非常频繁。

由于大语言模型的参数量巨大，许多大公司都需要训练数月，由此提出了各种资源消耗较小的训练方法，LoRA就是其中一种。

本文将详细介绍LoRA的原理，并使用PyTorch实现小模型的LoRA训练。

二、模型训练

现在大多数模型训练都是采用梯度下降算法。梯度下降算法可以分为下面4个步骤：

1. 正向传播计算损失值
2. 反向传播计算梯度
3. 利用梯度更新参数
4. 重复1、2、3的步骤，直到获取较小的损失

以线性模型为例，模型参数为W，输入输出为x、y，损失函数以均方误差为例。那么各个步骤的计算如下，首先是正向传播，对于线性模型来说就是做一个矩阵乘法：

在求出损失后，可以计算L对W的梯度，得到dW：

dW是一个矩阵，它会指向L上升最快的方向，但是我们的目的是让L下降，因此让W减去dW。为了调整更新的步伐，还会乘上一个学习率η，计算如下：

最后一直重复即刻。上述三个步骤的伪代码如下：

# 4、重复1、2、3
for i in range(10000):
    # 1、正向传播计算损失
    L = MSE(Wx, y)
    # 2、反向传播计算梯度
    dW = gradient(L, W)
    # 3、利用梯度更新参数
    W -= lr * dW

在更新完成后，得到新的参数W'。此时我们使用模型预测时，计算如下：

三、引入LoRA

我们可以来思考一下W和W'之间的关系。W通常指基础模型的参数，而W'是在基础模型的基础上，经过几次矩阵加减得到的。假设在训练的过程中更新了10次，每次的dW分别为dW1、dW2、....、dW10，那么完整的更新过程可以写为一次运算：

其中dW是一个形状与W'一致的矩阵。我们把-ηdW写成矩阵R，那么更新后的参数就是：

此时训练的过程就被简化为原矩阵加上另一个矩阵R。但是求解矩阵R并没有更简单，而且也没有节约资源，此时就引出LoRA了这一思想。

一个训练充分的矩阵，通常是满秩或者基本满足秩的，即矩阵中没有一列是多余的。在论文《Scaling Laws for Neural Language Model》中提出了数据集与参数大小之间的关系，满足该关系且训练良好，得到的模型是基本满秩的。在微调模型时，我们会选取一个底模，该底模就是基本满秩的。而更新矩阵R秩的情况是如何的呢？

我们假定R矩阵是一个低秩矩阵，低秩矩阵有许多重复的列，因此可以分解为两个更小的矩阵。假如W的形状为m×n，那么A的形状也是m×n，我们把矩阵R分解为AB（其中A形状为m×r，B形状为r×N），r通常会选取一个远小于m、n的值，如图所示：

将低秩矩阵分解为两个矩阵几点好处，首先是参数量明显减少。假设R矩阵的形状为100×100，那么R的参数量为10000。当我们选取秩为10时，此时矩阵A的形状为100×10，矩阵B的形状为10×100，此时参数量为2000，比R矩阵少了80%。

而且由于R是低秩矩阵，所以在训练充分的情况下，A和B矩阵可以达到R的效果。这里的矩阵AB就是我们常说的LoRA模型。

在引入LoRA后，我们的预测需要将x分别输入W和AB，此时预测的计算为：

在预测时会比原始模型稍慢，但是在大模型中基本感觉不到差异。

四、实战

为了把握各个细节，这里不使用大模型作为lora的实战，而是选择使用vgg19这种小型网络来训练lora模型。导入需要用到的模块：

import os  
import torch  
from torch import optim, nn  
from PIL import Image  
from torch.utils import data  
from torchvision import models  
from torchvision.transforms import transforms

4.1 数据集准备

这里使用vgg19在imagenet上的预训练权重作为底模，因此需要准备分类数据集。为了方便，这里只准备了一个类别，且只准备了5张图片，图片在项目下的data/goldfish下：

在imagenet中包含了goldfish类别，但是这里选取的是插画版的goldfish，经过测试，预训练模型不能将上述图片正确分类。我们的目的就是训练LoRA，让模型正确分类。

我们创建一个LoraDataset：

transform = transforms.Compose([  
    transforms.Resize(256),  
    transforms.CenterCrop(224),  
    transforms.ToTensor(),  
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  
])  
  
  
class LoraDataset(data.Dataset):  
    def __init__(self, data_path="datas"):  
        categories = models.VGG19_Weights.IMAGENET1K_V1.value.meta["categories"]  
        self.files = []  
        self.labels = []  
        for dir in os.listdir(data_path):  
            dirname = os.path.join(data_path, dir)  
            for file in os.listdir(dirname):  
                self.files.append(os.path.join(dirname, file))  
                self.labels.append(categories.index(dir))  
  
    def __getitem__(self, item):  
        image = Image.open(self.files[item]).convert("RGB")  
        label = torch.zeros(1000, dtype=torch.float64)  
        label[self.labels[item]] = 1.  
        return transform(image), label  
  
    def __len__(self):  
        return len(self.files)

4.2 创建LoRA模型

我们把LoRA封装成一个层，LoRA中只有两个需要训练的矩阵，LoRA的代码如下：

class Lora(nn.Module):  
    def __init__(self, m, n, rank=10):  
        super().__init__()  
        self.m = m  
        self.A = nn.Parameter(torch.randn(m, rank))  
        self.B = nn.Parameter(torch.zeros(rank, n))  
  
    def forward(self, inputs):  
        inputs = inputs.view(-1, self.m)  
        return torch.mm(torch.mm(inputs, self.A), self.B)

其中m是输入的大小，n是输出的大小，rank是秩的大小，我们可以设置一个较小的值。

在权重初始化时，我们把A用高斯噪声初始化，而B用0矩阵初始化，这样的目的是保证从底模开始训练。因为AB是0矩阵，所以初始状态下，LoRA不起作用。

4.3 设置超参数并训练

接下来就是训练了，这里和PyTorch常规训练代码基本一致，先看代码：

# 加载底模和lora  
vgg19 = models.vgg19(models.VGG19_Weights.IMAGENET1K_V1)  
for params in vgg19.parameters():  
    params.requires_grad = False  
vgg19.eval()  
lora = Lora(224 * 224 * 3, 1000)  
# 加载数据  
lora_loader = data.DataLoader(LoraDataset(), batch_size=batch_size, shuffle=True)  
# 加载优化器  
optimizer = optim.Adam(lora.parameters(), lr=lr)  
# 定义损失  
loss_fn = nn.CrossEntropyLoss()  
# 训练  
for epoch in range(epochs):  
    for image, label in lora_loader:  
        # 正向传播  
        pred = vgg19(image) + lora(image)  
        loss = loss_fn(pred, label)  
        # 反向传播  
        loss.backward()  
        # 更新参数  
        optimizer.step()  
        optimizer.zero_grad()  
        print(f"loss: {loss.item()}")

这里有两点需要注意，第一点是我们把vgg19的权重设置为不可训练，这和迁移学习很像，但其实是不一样的。

第二点则是正向传播时，我们使用了下面代码：

pred = vgg19(image) + lora(image)

4.4 测试

下面来简单测试一下：

# 测试  
for image, _ in lora_loader:  
    pred = vgg19(image) + lora(image)  
    idx = torch.argmax(pred, dim=1).item()  
    category = models.VGG19_Weights.IMAGENET1K_V1.value.meta["categories"][idx]  
    print(category)
torch.save(lora.state_dict(), 'lora.pth')

输出结果如下：

goldfish
goldfish
goldfish
goldfish
goldfish

基本预测正确了，不过这个测试结果并不能说明什么。最后我们保存了一个5M的LoRA模型，相比vgg19的几十M算是非常小了。

五、总结

LoRA是针对大模型的一种高效的训练方法，而本文则将LoRA使用在小型的分类网络中，旨在让读者更清晰认识LoRA的详细实现（同时也因为跑不动大模型）。限于数据量，对LoRA的精度效率等问题没有详细讨论，读者可以参考相关资料深入了解。

点赞 评论 收藏 目录

阅读量：1358

点赞量：0

收藏量:0