Stable Diffusion动态加载Lora过程中的实验、原理与说明

对比实验

显存占用情况

使用StableDiffusionXLPipeline.from_pretrained() 方法SDXL半精度加载显存占用约7G左右。

使用load_lora_weights()加载了5个Lora模型后显存提升到8G,平均每个Lora的大小在200M左右。 使用unload_lora_weights()后显存没有发生变化,还是8G,说明该方法不会清空已经加载到显存的Lora模型,但这时候再调用模型生成图片已经丢失Lora的效果了

推理耗时

Lora数量 耗时(秒)
0 15
1 20
2 24
... ...
7 45

这里使用的Lora平均每个的大小在200M左右,从上表不难发现单个Lora耗时约增加4秒左右。

代码分析与原理说明

1)加载Lora

通过调用load_lora_weights()来加载不同的Lora权重,这些权重的张量都会加载到显存中,但注意只有最后一次调用该方法的Lora才会生效load_lora_into_unet()load_lora_into_text_encoder(),之前的都会被覆盖掉。

python 复制代码
def load_lora_weights(
        self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], adapter_name=None, **kwargs
    ):
    
        ...
        # lora_state_dict 实际执行把tensor加载到显存中,同时返回2个字典记录所添加的lora的名称和配置信息
        state_dict, network_alphas = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
        
       
        # state_dict 和 network_alphas 是上面返回的2个参数
        # 所以每次load lora weight都会重新执行下面这2步
        # 因此加载后默认调用的lora就是最后load进来的lora
        self.load_lora_into_unet(
            state_dict,
            network_alphas=network_alphas,
            unet=getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet,
            low_cpu_mem_usage=low_cpu_mem_usage,
            adapter_name=adapter_name,
            _pipeline=self,
        )
        self.load_lora_into_text_encoder(
            state_dict,
            network_alphas=network_alphas,
            text_encoder=getattr(self, self.text_encoder_name)
            if not hasattr(self, "text_encoder")
            else self.text_encoder,
            lora_scale=self.lora_scale,
            low_cpu_mem_usage=low_cpu_mem_usage,
            adapter_name=adapter_name,
            _pipeline=self,
        )

用一张图来表示load_lora_weights()加载过程,蓝色表示生效的张量。

==因此如果不特别指定调用哪个Lora的话,默认输入的张量需要跟Base和最后一个Lora的权重分别相乘再叠加==

2)卸载Lora

卸载Lora时需要调用unload_lora_weights()方法,关键是把config删除,并不清显存,如代码中的del self.unet.peft_config操作。如果后续需要再调用Lora的话,则需要重头开始加载Lora权重

python 复制代码
  def unload_lora_weights(self):
        
        if not USE_PEFT_BACKEND:
            if version.parse(__version__) > version.parse("0.23"):
                logger.warn(
                    "You are using `unload_lora_weights` to disable and unload lora weights. If you want to iteratively enable and disable adapter weights,"
                    "you can use `pipe.enable_lora()` or `pipe.disable_lora()`. After installing the latest version of PEFT."
                )

            for _, module in self.unet.named_modules():
                if hasattr(module, "set_lora_layer"):
                    module.set_lora_layer(None)
        else:
            recurse_remove_peft_layers(self.unet)
            if hasattr(self.unet, "peft_config"):
            # 关键是把config删除,并不清显存
            # 因此执行unload_lora_weights 会丢失所有记载的lora的config信息,但不会释放显存
                del self.unet.peft_config

3)指定Lora

在拥有多个Lora加载后,如果需要指定某个或者某几个Lora,则需要用set_adapters()方法。该方法对Unet和Text Encoder都执行set_adapter()操作,间接调用了peft_utils 模块中的set_weights_and_activate_adapters() 方法为Unet或者Text Encoder的某些层指定Lora类型和权重。

python 复制代码
 def set_adapters(
        self,
        adapter_names: Union[List[str], str],
        adapter_weights: Optional[List[float]] = None,
    ):
        # 对unet和text encoder都执行set adapter操作。间接调用 set_weights_and_activate_adapters 方法
        self.unet.set_adapters(adapter_names, adapter_weights)

        # Handle the Text Encoder
        if hasattr(self, "text_encoder"):
            self.set_adapters_for_text_encoder(adapter_names, self.text_encoder, adapter_weights)
        if hasattr(self, "text_encoder_2"):
            self.set_adapters_for_text_encoder(adapter_names, self.text_encoder_2, adapter_weights)
python 复制代码
# src/diffusers/utils/peft_utils.py
def set_weights_and_activate_adapters(model, adapter_names, weights):
    from peft.tuners.tuners_utils import BaseTunerLayer

    # iterate over each adapter, make it active and set the corresponding scaling weight
    for adapter_name, weight in zip(adapter_names, weights):
        for module in model.modules():
            if isinstance(module, BaseTunerLayer):
                # For backward compatbility with previous PEFT versions
                if hasattr(module, "set_adapter"):
                    module.set_adapter(adapter_name)
                else:
                    module.active_adapter = adapter_name
                module.set_scale(adapter_name, weight)

如果已经执行过set_adapter(),而下次又需要使用不同的Lora,则重新执行set_adapter()选择需要的Lora即可

4)解除Lora

如果不需要Lora,想用Base模型直接生成图片,这时候可以通过disable_lora()方法来解除已经指定好的Lora。这时候再生成图片,输入的张量只会跟Base的矩阵张量相乘

python 复制代码
# src/diffusers/utils/peft_utils.py#L227
# disable_lora底层调用了该方法
# 会把Unet和Text Encoder的某些层注释掉Lora,因此达到解除Lora的效果
# 可以看到效果和指定Lora时刚好是相反的

def set_adapter_layers(model, enabled=True):
    from peft.tuners.tuners_utils import BaseTunerLayer

    for module in model.modules():
        if isinstance(module, BaseTunerLayer):
            # The recent version of PEFT needs to call `enable_adapters` instead
            if hasattr(module, "enable_adapters"):
            # false的时候注释掉 lora适配器
                module.enable_adapters(enabled=enabled)
            else:
                module.disable_adapters = not enabled

5)查看Lora

通过get_active_adapters()方法即可查看指定的Lora,即前面示例图中蓝色的Lora模块。

python 复制代码
    def get_active_adapters(self) -> List[str]:
     
        if not USE_PEFT_BACKEND:
            raise ValueError(
                "PEFT backend is required for this method. Please install the latest version of PEFT `pip install -U peft`"
            )

        from peft.tuners.tuners_utils import BaseTunerLayer

        active_adapters = []

        for module in self.unet.modules():
            if isinstance(module, BaseTunerLayer):
                active_adapters = module.active_adapters
                break

        return active_adapters

6)融合Lora

前面已经实验过,当有多个Lora调用时,推理时间几乎会线性增加,这是因为输入的张量除了与Base的模型相乘外,还需要与每个Lora相乘再叠加,因此推理时间会变长。一种解决方案是将需要的Lora与Base进行合并,即通过fuse_lora()来实现,这样推理时间可与Base单独使用时一致。下图中深蓝色表示Base的权重已经叠加了需要的Lora的权重,因此输入张量无需再经过Lora。但是缺点是==多次融合Lora后,Base将无法恢复,需要后续需要单独使用Base,只能重新加载Base!==

相关推荐
passer__jw7671 小时前
【LeetCode】【算法】3. 无重复字符的最长子串
算法·leetcode
passer__jw7671 小时前
【LeetCode】【算法】21. 合并两个有序链表
算法·leetcode·链表
sweetheart7-71 小时前
LeetCode22. 括号生成(2024冬季每日一题 2)
算法·深度优先·力扣·dfs·左右括号匹配
景鹤4 小时前
【算法】递归+回溯+剪枝:78.子集
算法·机器学习·剪枝
_OLi_4 小时前
力扣 LeetCode 704. 二分查找(Day1:数组)
算法·leetcode·职场和发展
丶Darling.5 小时前
Day40 | 动态规划 :完全背包应用 组合总和IV(类比爬楼梯)
c++·算法·动态规划·记忆化搜索·回溯
风影小子5 小时前
IO作业5
算法
奶味少女酱~5 小时前
常用的c++特性-->day02
开发语言·c++·算法
passer__jw7675 小时前
【LeetCode】【算法】11. 盛最多水的容器
算法·leetcode
诸葛悠闲5 小时前
C++ 面试问题集合
算法