对比实验
显存占用情况
使用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!==