Using LoRA for efficient fine-tuning: Fundamental principles — ROCm Blogs (amd.com)

[2106.09685] LoRA: Low-Rank Adaptation of Large Language Models (arxiv.org)

Parametrizations Tutorial — PyTorch Tutorials 2.3.0+cu121 documentation

大型语言模型（LLMs）的低秩适应（Low-Rank Adaptation，简称LoRA）用于解决微调大型语言模型时面临的挑战。像GPT和Llama这样的模型，拥有数十亿个参数，通常对于特定任务或领域的微调来说成本过高。LoRA保留了预训练模型的权重，并在每个模型块内部加入了可训练的层。这导致需要微调的参数数量显著减少，并大幅降低了GPU内存需求。LoRA的关键优势在于，它大幅减少了可训练参数的数量——有时高达10,000倍——从而显著降低了对GPU资源的需求。

为什么LoRA有效

当预训练的大型语言模型（LLMs）适应新任务时，它们具有较低的“内在维度”，这意味着数据可以在保留大部分关键信息或结构的同时，通过较低维度的空间进行有效表示或近似。我们可以通过将适应任务的新权重矩阵分解为较低维度（较小）的矩阵来实现这一点，而不会丢失太多重要信息。我们通过低秩近似来实现这一点。

矩阵的秩是一个值，它让你了解矩阵的复杂性。矩阵的低秩近似旨在尽可能接近地逼近原始矩阵，但具有较低的秩。较低秩的矩阵降低了计算复杂度，从而提高了矩阵乘法的效率。低秩分解是指通过导出A的低秩近似来有效逼近矩阵A的过程。奇异值分解（SVD）是低秩分解的一种常见方法。

假设W表示给定神经网络层中的权重矩阵，假设ΔW是经过完整微调后W的权重更新。然后，我们可以将权重更新矩阵ΔW分解为两个较小的矩阵：ΔW = WA*WB，其中WA是A×r维矩阵，WB是r×B维矩阵。在这里，我们保持原始权重W不变，只训练新的矩阵WA和WB。这总结了LoRA方法，如下图所示。

LoRA的好处

1. 降低资源消耗：微调深度学习模型通常需要大量的计算资源，这可能既昂贵又耗时。LoRA在保持高性能的同时，降低了对资源的需求。

2. 更快的迭代：LoRA能够实现快速迭代，使得尝试不同的微调任务和快速适应模型变得更加容易。

3. 改进的迁移学习：LoRA增强了迁移学习的效果，因为带有LoRA适配器的模型可以用更少的数据进行微调。这在标签数据稀缺的情况下尤其有价值。

4. 广泛的应用性：LoRA具有通用性，可以应用于各种领域，包括自然语言处理、计算机视觉和语音识别等。

5. 更低的碳足迹：通过降低计算需求，LoRA有助于实现更绿色、更可持续的深度学习方法。

使用LoRA技术训练神经网络

我们的目标是使用LoRA技术训练一个神经网络，用于MNIST手写数字数据库的分类任务，并随后对该网络进行微调，以提升其在初始表现不佳的类别上的性能。

硬件要求

• AMD Instinct GPU

软件环境

• ROCm：ROCm是针对AMD GPU优化的开源机器学习平台。
• PyTorch：PyTorch是广泛使用的深度学习框架，支持动态计算图。
• tqdm：Python库，用于显示进度条，方便观察训练过程。

以下是一个简化的步骤指南，说明如何使用LoRA技术训练神经网络（注意：具体代码细节可能需要根据您的环境和PyTorch版本进行调整）：

1. 准备数据集：
   首先，您需要准备MNIST数据集。PyTorch提供了torchvision.datasets.MNIST来方便加载这个数据集。

2. 定义模型：
   选择一个预训练的神经网络模型，例如ResNet、Transformer等。为了简单起见，您可以选择一个轻量级的卷积神经网络（CNN）。

3. 添加LoRA层：
   在模型的每个权重矩阵上添加LoRA层。这通常涉及在模型内部插入额外的可训练权重矩阵（如WA和WB），它们将用于生成对原始权重矩阵的低秩更新。

4. 初始化模型：
   加载预训练的模型权重，并将LoRA层的权重初始化为零或小的随机值。

5. 设置训练循环：
   使用PyTorch的数据加载器（DataLoader）和优化器（如SGD或Adam）来设置训练循环。确保在训练循环中同时更新原始模型的权重和LoRA层的权重。

6. 训练模型：
   运行训练循环，通过反向传播和优化器步骤来更新权重。由于LoRA层的存在，大多数权重更新将仅应用于LoRA层，而不是整个模型。

7. 评估模型：
   在验证集上评估模型的性能，并根据需要调整超参数或继续训练。

8. 微调模型：
   如果模型在特定类别上表现不佳，可以使用LoRA技术仅对该类别进行微调。这通常涉及重新训练LoRA层（同时保持原始模型权重不变），使用仅针对该类别的数据。

9. 测试模型：
   在测试集上测试微调后的模型，以评估其性能改进。

开始

1. 首先，我们需要导入一些必要的包。

import torch
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import torch.nn as nn
from tqdm import tqdm

2. 设定随机数生成的种子，以确保模型的行为在每次运行时都是确定的。

# 设置随机种子以保证实验的可复现性
_ = torch.manual_seed(0)

我们通常不会将变量名 _ 作为赋值语句的结果，除非我们故意忽略该值。在这里，它仅用于表示我们不关心 torch.manual_seed(0) 的返回值，并使其确定性。

在训练神经网络时，我们通常希望结果是可重复的，这意味着如果我们用相同的初始参数和相同的训练数据重新运行模型，我们应该得到相同的结果。通过设置随机数生成的种子（使用 torch.manual_seed(0)），我们可以确保PyTorch在生成随机数（例如，在初始化权重或选择随机数据批次时）时使用相同的序列，从而使模型训练具有确定性。这在调试和比较不同模型或超参数时特别有用。

3. 加载数据集。

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])

# 加载MNIST数据集
mnist_trainset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
# 为训练创建数据加载器
train_loader = torch.utils.data.DataLoader(mnist_trainset, batch_size=10, shuffle=True)

# 加载MNIST测试集
mnist_testset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
test_loader = torch.utils.data.DataLoader(mnist_testset, batch_size=10, shuffle=True)

# 定义设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

4. 创建用于分类数字的神经网络（我们使用了更复杂的代码以更好地展示LoRA）。

# 创建一个过度昂贵的神经网络来分类MNIST数字
# 不关心效率
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)

这段代码首先定义了一个数据预处理流程，用于将MNIST数据集中的图像转换为张量并进行归一化。然后，它加载了MNIST训练集和测试集，并为这两个数据集创建了数据加载器。接下来，它定义了一个设备，用于确定是在GPU上还是在CPU上运行模型。最后，它定义了一个名为RichBoyNet的神经网络类，该类包含三个全连接层和一个ReLU激活函数。最后，它创建了这个网络的实例，并将其移动到指定的设备上。

5. 对网络进行一轮训练，以模拟在数据上的完整预训练过程。在AMD Instinct GPU上，此过程只需数秒。

这段代码定义了一个名为train的函数，用于训练一个神经网络模型（在这里称为net）一个或多个epoch（周期）。在每个epoch中，该函数会遍历训练数据集（由train_loader提供），并更新模型的权重以最小化预测输出和实际标签之间的交叉熵损失。

# 定义一个函数来训练网络  
def train(train_loader, net, epochs=5, total_iterations_limit=None):  
    # 使用交叉熵损失函数  
    cross_el = nn.CrossEntropyLoss()  
    # 使用Adam优化器来更新网络参数，学习率设置为0.001  
    optimizer = torch.optim.Adam(net.parameters(), lr=0.001)  
  
    total_iterations = 0  # 初始化总迭代次数为0  
  
    # 对于指定的epochs数量，进行循环  
    for epoch in range(epochs):  
        net.train()  # 设置网络为训练模式  
  
        loss_sum = 0  # 初始化损失总和为0  
        num_iterations = 0  # 初始化迭代次数为0  
  
        # 使用tqdm库包装train_loader，使其具有进度条功能，并显示当前epoch信息  
        data_iterator = tqdm(train_loader, desc=f'Epoch {epoch+1}')  
        # 如果指定了总的迭代次数限制，则更新tqdm的total值  
        if total_iterations_limit is not None:  
            data_iterator.total = total_iterations_limit  
  
        # 遍历训练数据  
        for data in data_iterator:  
            num_iterations += 1  # 迭代次数加1  
            total_iterations += 1  # 总迭代次数加1  
            x, y = data  # 解包数据（输入x和标签y）  
            x = x.to(device)  # 将输入数据移动到指定的设备（如GPU）  
            y = y.to(device)  # 将标签数据移动到指定的设备  
  
            optimizer.zero_grad()  # 清零优化器的梯度  
            # 将输入数据展平（假设每个输入图像是28x28的），然后传递给网络  
            output = net(x.view(-1, 28*28))  
            # 计算预测输出和实际标签之间的交叉熵损失  
            loss = cross_el(output, y)  
            loss_sum += loss.item()  # 累加损失值  
            avg_loss = loss_sum / num_iterations  # 计算平均损失  
            # 更新tqdm的进度条，显示当前平均损失  
            data_iterator.set_postfix(loss=avg_loss)  
  
            # 执行反向传播以计算梯度  
            loss.backward()  
            # 使用优化器更新网络参数  
            optimizer.step()  
  
            # 如果指定了总的迭代次数限制，并且已经达到该限制，则退出训练  
            if total_iterations_limit is not None and total_iterations >= total_iterations_limit:  
                return  
  
# 调用train函数，只训练一个epoch  
train(train_loader, net, epochs=1)

在这段代码中，device应该是一个预定义的变量，表示要使用的设备（CPU或GPU）。在实际使用中，你需要先定义device，例如device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')，然后才能将数据移动到该设备上。此外，tqdm是一个用于显示进度条的Python库，如果你没有安装它，需要先使用pip install tqdm进行安装。

定义一个函数train用于对指定的数据加载器train_loader和网络net进行训练，可以选择训练的周期数epochs，默认为5，以及总的迭代次数限制total_iterations_limit。

在函数中，首先初始化交叉熵损失函数CrossEntropyLoss和Adam优化器。然后，循环进行每个周期的训练：

• 设置网络为训练模式。
• 初始化损失总和loss_sum和迭代次数num_iterations。
• 使用tqdm包装训练数据加载器，以便显示进度。
• 迭代数据加载器中的每个数据样本，执行以下操作：
  • 更新迭代次数。
  • 将输入数据x和标签y转换为适合GPU的格式。
  • 清除优化器的梯度。
  • 通过前向传播计算网络的输出。
  • 计算损失。
  • 累加损失并计算平均损失。
  • 将平均损失显示在进度条上。
  • 执行反向传播。
  • 更新网络权重。
  • 如果达到总迭代次数限制，提前结束训练。

最后，调用train函数，传入训练数据加载器、网络对象，并指定训练周期为1。

[!提示] 保留原始权重的副本（克隆它们），以便在微调后查看原始权重是否被更改。

original_weights = {}  
for name, param in net.named_parameters():  
    original_weights[name] = param.clone().detach()

在微调神经网络之前，我们首先遍历网络中的每一个参数（权重和偏置），并将它们保存在一个名为original_weights的字典中。这样，我们就可以在微调过程之后，通过比较original_weights中的权重和微调后网络中的权重，来查看原始权重是否已经被更改。

param.clone().detach()确保了我们在保存权重时创建了一个与原始权重完全相同的独立副本，这样即使在微调过程中原始权重发生更改，这个副本也不会受到影响。其中，clone()方法用于复制参数，detach()方法用于断开计算图，确保这个副本不参与任何梯度计算。

微调

1. 选择一个数字进行微调。预训练的网络在数字9上的表现不佳，所以我们将对这个数字进行微调。

def test():  
    correct = 0  
    total = 0  
  
    wrong_counts = [0 for i in range(10)]  # 初始化一个长度为10的列表，用于记录每个数字的错误次数  
  
    with torch.no_grad():  # 禁用梯度计算，因为我们只进行前向传播来评估模型  
        for data in tqdm(test_loader, desc='Testing'):  # tqdm是一个用于显示进度的工具  
            x, y = data  
            x = x.to(device)  # 将输入数据移动到指定的设备（如GPU）  
            y = y.to(device)  # 将标签数据移动到指定的设备  
            output = net(x.view(-1, 784))  # 假设每个图像是28x28的，因此将其展平为784个特征  
  
            # 遍历每个输出和对应的标签  
            for idx, i in enumerate(output):  
                if torch.argmax(i) == y[idx]:  # 如果预测的最大概率索引与真实标签相同  
                    correct += 1  # 正确预测数加1  
                else:  
                    wrong_counts[y[idx]] += 1  # 将对应标签的错误次数加1  
                total += 1  # 总测试样本数加1  
  
    print(f'Accuracy: {round(correct/total, 3)}')  # 打印准确率  
  
    # 打印每个数字的错误次数  
    for i in range(len(wrong_counts)):  
        print(f'数字 {i} 的错误次数: {wrong_counts[i]}')  
  
# 调用测试函数  
test()

这段代码首先定义了一个test函数，用于在测试集上评估模型的性能。它计算了模型预测正确的样本数量，并统计了每个数字被错误预测的次数。然后，它输出了模型的总准确率和每个数字的错误次数。最后，通过调用test()函数来执行测试。

定义一个测试函数test，用于评估网络在MNIST测试集上的性能。首先初始化正确预测的数量correct和总预测数量total，以及一个用于统计每个数字被错误识别次数的列表wrong_counts。

在不计算梯度的上下文中（torch.no_grad()），遍历测试数据加载器test_loader，并使用进度条显示测试进度。对于每个数据样本：

• 将输入数据x和标签y转换为适合GPU的格式。
• 通过前向传播计算网络的输出。
• 遍历输出中每个预测结果，如果预测的类别与真实标签相同，则增加正确预测的数量；如果不同，则增加该数字的错误计数。
• 更新总预测数量。

测试完成后，打印出整体的准确率，并遍历wrong_counts列表，打印出每个数字的错误计数。

最后，调用test函数执行测试过程。

2. 在引入LoRA矩阵之前，可视化原始网络中参数的数量。

# 打印网络权重矩阵的大小
# 保存总参数数量的计数
total_parameters_original = 0
for index, layer in enumerate([net.linear1, net.linear2, net.linear3]):
    total_parameters_original += layer.weight.nelement() + layer.bias.nelement()
    print(f'Layer {index+1}: W: {layer.weight.shape} + B: {layer.bias.shape}')
print(f'Total number of parameters: {total_parameters_original:,}')

上面的代码段通过遍历网络中的线性层（假设net.linear1、net.linear2和net.linear3是网络中的线性层），计算了每个层的权重（W）和偏置（B）的元素数量，并将它们累加到total_parameters_original变量中。最后，它打印出了每个层的权重和偏置的形状以及网络中的总参数数量。

3. 定义LoRA参数化。

LoRA（Low-Rank Adaptation）参数化是一种用于微调大型神经网络模型的技术，特别是那些已经经过预训练的模型。LoRA通过在原始权重矩阵上添加一个低秩更新项来实现参数的高效微调，而不是直接更新整个权重矩阵。

class LoRAParametrization(nn.Module):  
    def __init__(self, features_in, features_out, rank=1, alpha=1, device='cpu'):  
        super().__init__()  
        # 论文第4.1节：  
        # 我们使用随机高斯初始化A，并将B初始化为零，所以ΔW = BA在训练开始时为零  
        self.lora_A = nn.Parameter(torch.zeros((rank, features_out)).to(device))  # 初始化为零的低秩矩阵A  
        self.lora_B = nn.Parameter(torch.zeros((features_in, rank)).to(device))  # 初始化为零的低秩矩阵B  
        nn.init.normal_(self.lora_A, mean=0, std=1)  # 对A进行标准正态分布初始化  
  
        # 论文第4.1节：  
        # 我们对ΔWx进行α/r的缩放，其中α是一个与r无关的常数。  
        # 当使用Adam进行优化时，调整α大致相当于调整学习率，如果我们适当地缩放初始化。  
        # 因此，我们简单地将α设置为我们尝试的第一个r，并不调整它。  
        # 这种缩放有助于在改变r时减少重新调整超参数的需要。  
        self.scale = alpha / rank  # 缩放因子  
        self.enabled = True  # 标志位，表示是否启用LoRA更新  
  
    def forward(self, original_weights):  
        if self.enabled:  
            # 返回 W + (B*A)*scale  
            # 这里B*A是一个低秩矩阵，用于更新原始权重W  
            return original_weights + torch.matmul(self.lora_B, self.lora_A).view(original_weights.shape) * self.scale  
        else:  
            # 如果未启用LoRA，则直接返回原始权重  
            return original_weights

这个类定义了一个LoRA参数化模块，它可以在原始权重上添加一个低秩更新项。这个更新项由两个低秩矩阵lora_A和lora_B的乘积以及一个缩放因子scale组成。在训练过程中，这个模块会优化这两个低秩矩阵，而不是整个权重矩阵，从而大大减少了需要更新的参数数量，提高了微调的效率和效果。

定义一个名为LoRAParametrization的类，它是一个神经网络模块。

在初始化函数__init__中：

• 调用父类nn.Module的初始化函数。

• 根据论文的4.1节，对lora_A使用随机高斯初始化，对lora_B使用零初始化，这样训练开始时∆W = BA为零。

• lora_A和lora_B被定义为神经网络的参数，分别初始化为大小为(rank, features_out)和(features_in, rank)的零矩阵，并指定设备device。

• 使用nn.init.normal_函数对lora_A进行正态分布初始化，均值为0，标准差为1。

• 根据论文的4.1节，∆Wx按α/r的比例进行缩放，其中α是r中的一个常数。

• 当使用Adam优化器时，调整α大致上与调整学习率相同，如果我们适当地缩放初始化。

• 因此，我们简单地将α设置为我们尝试的第一个r的值，并且不进行调整。

• 这种缩放有助于减少我们在变化r时需要重新调整超参数的需求。

• 计算缩放因子self.scale = alpha / rank。

• 设置一个标志self.enabled为True，表示LoRA参数化被启用。

在前向传播函数forward中：

• 如果self.enabled为True，则计算W + (B*A)*scale，即将原始权重original_weights与其对应的lora_B和lora_A的矩阵乘法结果进行缩放后相加。
• 如果self.enabled为False，则直接返回原始权重original_weights。

4. 将参数化添加到我们的网络中。可以在PyTorch.org上了解更多关于PyTorch参数化的信息。

import torch.nn.utils.parametrize as parametrize  
  
# 定义一个函数来将LoRA参数化应用到线性层的权重矩阵上  
def linear_layer_parameterization(layer, device, rank=1, lora_alpha=1):  
    # 仅对权重矩阵添加参数化，忽略偏置项  
  
    # 根据论文第4.2节：  
    # 我们的研究仅限于为下游任务调整注意力权重，并冻结MLP模块（因此在下游任务中它们不会被训练）  
    # [...]  
    # 我们将[...]和偏置项的实证研究留给未来的工作。  
  
    features_in, features_out = layer.weight.shape  
    return LoRAParametrization(  
        features_in, features_out, rank=rank, alpha=lora_alpha, device=device  
    )  
  
# 使用parametrize.register_parametrization将LoRA参数化应用到指定的网络层  
# 这里我们为net.linear1, net.linear2, net.linear3的权重添加了LoRA参数化  
parametrize.register_parametrization(  
    net.linear1, "weight", linear_layer_parameterization(net.linear1, device="your_device_here")  
)  
parametrize.register_parametrization(  
    net.linear2, "weight", linear_layer_parameterization(net.linear2, device="your_device_here")  
)  
parametrize.register_parametrization(  
    net.linear3, "weight", linear_layer_parameterization(net.linear3, device="your_device_here")  
)  
  
# 注意：上面的代码中 "your_device_here" 需要替换为你的实际设备名，比如 "cuda:0" 或 "cpu"  
  
# 定义一个函数来启用或禁用LoRA参数化  
def enable_disable_lora(enabled=True):  
    for layer in [net.linear1, net.linear2, net.linear3]:  
        # 通过访问layer.parametrizations["weight"][0]来访问并修改LoRA参数化的enabled属性  
        layer.parametrizations["weight"][0].enabled = enabled

在上面的代码中，我们首先定义了一个函数linear_layer_parameterization，它接受一个线性层、设备名称、秩（rank）和LoRA的alpha参数作为输入，并返回一个LoRAParametrization对象。然后，我们使用parametrize.register_parametrization函数将这个参数化应用到网络的特定层的权重上。最后，我们定义了一个enable_disable_lora函数，它允许我们动态地启用或禁用这些层上的LoRA参数化。

请注意，你需要将"your_device_here"替换为适合你实际情况的设备名称，比如"cuda:0"（如果你的GPU可用并且你想在GPU上运行）或"cpu"（如果你想在CPU上运行）。

导入torch.nn.utils.parametrize模块，这个模块提供了参数化网络权重的工具。

定义一个函数linear_layer_parameterization，用于生成LoRA参数化对象。此函数接收一个层对象layer、设备device、秩rank和LoRA参数lora_alpha作为参数，然后返回一个LoRAParametrization实例。

• 根据论文的4.2节，本研究仅限于仅适应下游任务的注意力权重，并将多层感知器（MLP）模块冻结，这样做既简单又节省参数。
• 函数中获取层的权重矩阵的形状features_in和features_out。
• 返回一个LoRAParametrization实例，该实例将用于参数化指定层的权重矩阵。

使用parametrize.register_parametrization函数为网络中的线性层（net.linear1、net.linear2、net.linear3）的权重注册LoRA参数化。

定义一个函数enable_disable_lora，用于启用或禁用LoRA参数化。

• 此函数接收一个参数enabled，用于指定是否启用LoRA参数化，默认为True启用。
• 遍历网络中的线性层，通过layer.parametrizations["weight"][0].enabled设置LoRA参数化的启用状态。

这样，我们就可以在网络的每个线性层上应用LoRA参数化，并根据需要启用或禁用它。

5. 显示由LoRA添加的参数数量。

# 初始化LoRA参数和非LoRA参数的总数  
total_parameters_lora = 0  
total_parameters_non_lora = 0  
  
# 遍历网络中的线性层  
for index, layer in enumerate([net.linear1, net.linear2, net.linear3]):  
    # 计算LoRA参数（lora_A和lora_B）的数量，并累加到total_parameters_lora  
    total_parameters_lora += layer.parametrizations["weight"][0].lora_A.nelement() + layer.parametrizations["weight"][0].lora_B.nelement()  
      
    # 计算非LoRA参数（权重和偏置）的数量，并累加到total_parameters_non_lora  
    total_parameters_non_lora += layer.weight.nelement() + layer.bias.nelement()  
      
    # 打印每层的权重、偏置以及LoRA参数A和B的形状  
    print(  
        f'层 {index+1}: W: {layer.weight.shape} + B: {layer.bias.shape} + Lora_A: {layer.parametrizations["weight"][0].lora_A.shape} + Lora_B: {layer.parametrizations["weight"][0].lora_B.shape}'  
    )  
  
# 验证非LoRA参数的总数是否与原始网络的参数总数匹配  
# 注意：total_parameters_original需要在代码的其他部分定义  
assert total_parameters_non_lora == total_parameters_original  
  
# 打印原始网络的参数总数  
print(f'原始参数总数: {total_parameters_non_lora:,}')  
  
# 打印原始参数和LoRA参数的总数  
print(f'原始参数 + LoRA参数的总数: {total_parameters_lora + total_parameters_non_lora:,}')  
  
# 打印LoRA引入的参数数量  
print(f'LoRA引入的参数数量: {total_parameters_lora:,}')  
  
# 计算LoRA参数相对于原始参数的增量百分比  
parameters_increment = (total_parameters_lora / total_parameters_non_lora) * 100  
print(f'参数增量: {parameters_increment:.3f}%')

计算并显示通过LoRA技术添加到网络中的参数数量。

• 初始化total_parameters_lora和total_parameters_non_lora变量，分别用于存储LoRA参数和非LoRA参数的总数。
• 遍历网络中的每个线性层（net.linear1, net.linear2, net.linear3）：
  • 将每个层的LoRA参数（lora_A和lora_B）的元素数量加到total_parameters_lora。
  • 将每个层的权重和偏置的元素数量加到total_parameters_non_lora。
  • 打印每层的权重、偏置、LoRA权重矩阵A和B的形状。
• 通过断言检查非LoRA参数的数量是否与原始网络中的参数数量相同，确保没有计算错误。
• 打印原始网络的总参数数量、添加LoRA后的总参数数量，以及LoRA引入的参数数量。
• 计算参数增加的比例，并以百分比的形式打印出来，保留三位小数。

这段代码计算并打印出在应用LoRA参数化后，每个线性层增加的LoRA参数（lora_A和lora_B）的总数量，同时也计算并验证了非LoRA参数（原始权重和偏置项）的总数量与未应用LoRA之前的总参数量一致。最后，它报告了总体参数量的增加情况，包括原始参数和LoRA参数的总和，以及LoRA单独引入的参数数量及其相对于非LoRA参数的百分比增长。

6. 冻结原始网络的所有参数，仅微调LoRA引入的参数。然后，针对数字9微调模型100个批次。

# 冻结非LoRA参数
for name, param in net.named_parameters():
    if 'lora' not in name:
        print(f'Freezing non-LoRA parameter {name}')
        param.requires_grad = False

# 重新加载MNIST数据集，仅保留数字9
mnist_trainset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
exclude_indices = mnist_trainset.targets == 9
mnist_trainset.data = mnist_trainset.data[exclude_indices]
mnist_trainset.targets = mnist_trainset.targets[exclude_indices]
# 为训练创建一个数据加载器
train_loader = torch.utils.data.DataLoader(mnist_trainset, batch_size=10, shuffle=True)

# 仅使用LoRA在数字9上训练网络，并且只训练100个批次（希望它能提高数字9的性能）
train(train_loader, net, epochs=1, total_iterations_limit=100)

在这个设置中，我们仅对LoRA引入的参数进行微调，以期望改善模型在数字9上的性能，同时保持原始网络的其他参数不变。我们重新加载了MNIST数据集，并过滤出所有不是数字9的样本，以便我们的模型能够专注于学习数字9的特征。然后，我们创建了一个数据加载器，并使用一个自定义的train函数（需要用户自行定义）来训练网络，限制训练迭代次数为100次。

7. 验证微调没有改变原始权重（仅使用LoRA引入的权重）。

# 检查冻结的参数在微调后仍然不变
assert torch.all(net.linear1.parametrizations.weight.original == original_weights['linear1.weight'])
assert torch.all(net.linear2.parametrizations.weight.original == original_weights['linear2.weight'])
assert torch.all(net.linear3.parametrizations.weight.original == original_weights['linear3.weight'])

enable_disable_lora(enabled=True)
# 现在让我们以net.linear1层为例，检查LoRA是否正确应用到模型中，如LoRAParametrization.forward()中定义
# 新的linear1.weight是通过我们的LoRA参数化的"forward"函数获得的
# 原始权重已经移动到了net.linear1.parametrizations.weight.original
# 更多信息在这里：https://pytorch.org/tutorials/intermediate/parametrizations.html#inspecting-a-parametrized-module
assert torch.equal(net.linear1.weight, net.linear1.parametrizations.weight.original + (net.linear1.parametrizations.weight[0].lora_B @ net.linear1.parametrizations.weight[0].lora_A) * net.linear1.parametrizations.weight[0].scale)

enable_disable_lora(enabled=False)
# 如果我们禁用LoRA，linear1.weight就是原始的那个
assert torch.equal(net.linear1.weight, original_weights['linear1.weight'])

这段代码首先验证了在微调后，所有原始的权重确实没有被改变。接着，启用LoRA，通过断言检查net.linear1层的权重是否如预期那样由原始权重加上LoRA参数化后的调整构成。最后，当禁用LoRA时，再次验证net.linear1.weight确实恢复为原始权重，确保了LoRA的开关功能正常工作。

8. 启用LoRA进行测试。启用LoRA功能，然后进行测试。数字9应该被分类得更好。

# 使用LoRA启用进行测试  
enable_disable_lora(enabled=True)  
test()  # 假设test()函数用于测试网络并输出相关结果

禁用LoRA进行测试。禁用LoRA功能，然后再次进行测试。准确率和错误计数应该与原始网络相同。

# 使用LoRA禁用进行测试  
enable_disable_lora(enabled=False)  
test()  # 假设test()函数用于测试网络并输出相关结果

[!注意] 您可能会观察到微调对其他标签的准确率产生了影响。这是预期的，因为我们的微调是专门针对数字9进行的。

完整代码

import torch
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import torch.nn as nn
from tqdm import tqdm

# Make torch deterministic
_ = torch.manual_seed(0)

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])

# Load the MNIST data set
mnist_trainset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
# Create a dataloader for the training
train_loader = torch.utils.data.DataLoader(mnist_trainset, batch_size=10, shuffle=True)

# Load the MNIST test set
mnist_testset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
test_loader = torch.utils.data.DataLoader(mnist_testset, batch_size=10, shuffle=True)

# Define the device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# Create an overly expensive neural network to classify MNIST digits
# Daddy got money, so I don't care about efficiency
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)

def train(train_loader, net, epochs=5, total_iterations_limit=None):
    cross_el = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(net.parameters(), lr=0.001)

    total_iterations = 0

    for epoch in range(epochs):
        net.train()

        loss_sum = 0
        num_iterations = 0

        data_iterator = tqdm(train_loader, desc=f'Epoch {epoch+1}')
        if total_iterations_limit is not None:
            data_iterator.total = total_iterations_limit
        for data in data_iterator:
            num_iterations += 1
            total_iterations += 1
            x, y = data
            x = x.to(device)
            y = y.to(device)
            optimizer.zero_grad()
            output = net(x.view(-1, 28*28))
            loss = cross_el(output, y)
            loss_sum += loss.item()
            avg_loss = loss_sum / num_iterations
            data_iterator.set_postfix(loss=avg_loss)
            loss.backward()
            optimizer.step()

            if total_iterations_limit is not None and total_iterations >= total_iterations_limit:
                return

train(train_loader, net, epochs=1)

original_weights = {}
for name, param in net.named_parameters():
    original_weights[name] = param.clone().detach()

def test():
    correct = 0
    total = 0

    wrong_counts = [0 for i in range(10)]

    with torch.no_grad():
        for data in tqdm(test_loader, desc='Testing'):
            x, y = data
            x = x.to(device)
            y = y.to(device)
            output = net(x.view(-1, 784))
            for idx, i in enumerate(output):
                if torch.argmax(i) == y[idx]:
                    correct +=1
                else:
                    wrong_counts[y[idx]] +=1
                total +=1
    print(f'Accuracy: {round(correct/total, 3)}')
    for i in range(len(wrong_counts)):
        print(f'wrong counts for the digit {i}: {wrong_counts[i]}')

test()

# Print the size of the weights matrices of the network
# Save the count of the total number of parameters
total_parameters_original = 0
for index, layer in enumerate([net.linear1, net.linear2, net.linear3]):
    total_parameters_original += layer.weight.nelement() + layer.bias.nelement()
    print(f'Layer {index+1}: W: {layer.weight.shape} + B: {layer.bias.shape}')
print(f'Total number of parameters: {total_parameters_original:,}')

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 training
        self.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 / rank
        self.enabled = True

    def forward(self, original_weights):
        if self.enabled:
            # Return W + (B*A)*scale
            return original_weights + torch.matmul(self.lora_B, self.lora_A).view(original_weights.shape) * self.scale
        else:
            return original_weights
        
import torch.nn.utils.parametrize as parametrize

def linear_layer_parameterization(layer, device, rank=1, lora_alpha=1):
    # Only add the parameterization to the weight matrix, ignore the Bias

    # From section 4.2 of the paper:
    # We limit our study to only adapting the attention weights for downstream tasks and freeze the
    # MLP modules (so they are not trained in downstream tasks) both for simplicity and
    # parameter-efficiency.
    # [...]
    # We leave the empirical investigation of [...], and biases to a future work.

    features_in, features_out = layer.weight.shape
    return LoRAParametrization(
        features_in, features_out, rank=rank, alpha=lora_alpha, device=device
    )

parametrize.register_parametrization(
    net.linear1, "weight", linear_layer_parameterization(net.linear1, device)
)
parametrize.register_parametrization(
    net.linear2, "weight", linear_layer_parameterization(net.linear2, device)
)
parametrize.register_parametrization(
    net.linear3, "weight", linear_layer_parameterization(net.linear3, device)
)


def enable_disable_lora(enabled=True):
    for layer in [net.linear1, net.linear2, net.linear3]:
        layer.parametrizations["weight"][0].enabled = enabled

total_parameters_lora = 0
total_parameters_non_lora = 0
for index, layer in enumerate([net.linear1, net.linear2, net.linear3]):
    total_parameters_lora += layer.parametrizations["weight"][0].lora_A.nelement() + layer.parametrizations["weight"][0].lora_B.nelement()
    total_parameters_non_lora += layer.weight.nelement() + layer.bias.nelement()
    print(
        f'Layer {index+1}: W: {layer.weight.shape} + B: {layer.bias.shape} + Lora_A: {layer.parametrizations["weight"][0].lora_A.shape} + Lora_B: {layer.parametrizations["weight"][0].lora_B.shape}'
    )
# The non-LoRA parameters count must match the original network
assert total_parameters_non_lora == total_parameters_original
print(f'Total number of parameters (original): {total_parameters_non_lora:,}')
print(f'Total number of parameters (original + LoRA): {total_parameters_lora + total_parameters_non_lora:,}')
print(f'Parameters introduced by LoRA: {total_parameters_lora:,}')
parameters_increment = (total_parameters_lora / total_parameters_non_lora) * 100
print(f'Parameters increment: {parameters_increment:.3f}%')

# Freeze the non-Lora parameters
for name, param in net.named_parameters():
    if 'lora' not in name:
        print(f'Freezing non-LoRA parameter {name}')
        param.requires_grad = False

# Load the MNIST data set again, by keeping only the digit 9
mnist_trainset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
exclude_indices = mnist_trainset.targets == 9
mnist_trainset.data = mnist_trainset.data[exclude_indices]
mnist_trainset.targets = mnist_trainset.targets[exclude_indices]
# Create a dataloader for the training
train_loader = torch.utils.data.DataLoader(mnist_trainset, batch_size=10, shuffle=True)

# Train the network with LoRA only on the digit 9 and only for 100 batches (hoping that it would
# improve the performance on the digit 9)
train(train_loader, net, epochs=1, total_iterations_limit=100)

# Check that the frozen parameters are still unchanged by the finetuning
assert torch.all(net.linear1.parametrizations.weight.original == original_weights['linear1.weight'])
assert torch.all(net.linear2.parametrizations.weight.original == original_weights['linear2.weight'])
assert torch.all(net.linear3.parametrizations.weight.original == original_weights['linear3.weight'])

enable_disable_lora(enabled=True)
# Now let's use layer of net.linear1 as an example to check if the Lora is applied to the model
# correctly as defined in the LoRAParametrization.forward()
# The new linear1.weight is obtained by the "forward" function of our LoRA parametrization
# The original weights have been moved to net.linear1.parametrizations.weight.original
# More info here: https://pytorch.org/tutorials/intermediate/parametrizations.html#inspecting-a-parametrized-module
assert torch.equal(net.linear1.weight, net.linear1.parametrizations.weight.original + (net.linear1.parametrizations.weight[0].lora_B @ net.linear1.parametrizations.weight[0].lora_A) * net.linear1.parametrizations.weight[0].scale)

enable_disable_lora(enabled=False)
# If we disable LoRA, the linear1.weight is the original one
assert torch.equal(net.linear1.weight, original_weights['linear1.weight'])

# Test with LoRA enabled
enable_disable_lora(enabled=True)
test()

# Test with LoRA disabled
enable_disable_lora(enabled=False)
test()