图像生成管线：LoRA微调

import torchimport torch.nn as nnimport math class LoRALayer(nn.Module):    def __init__(self, original_layer, rank=16, alpha=32):        super().__init__()        self.original = original_layer        self.original.weight.requires_grad = False        in_features = original_layer.in_features        out_features = original_layer.out_features        self.rank = rank        self.alpha = alpha        self.scaling = alpha / rank        self.lora_A = nn.Parameter(torch.empty(in_features, rank))        self.lora_B = nn.Parameter(torch.zeros(rank, out_features))        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))    def forward(self, x):        base_output = self.original(x)        lora_output = x @ self.lora_A @ self.lora_B * self.scaling        return base_output + lora_output    def merge_weights(self):        merged_weight = self.original.weight.data + (self.lora_A @ self.lora_B * self.scaling).T        self.original.weight.data.copy_(merged_weight)        return self.original class LoRAManager:    def __init__(self):        self.loaded_loras = {}    def register_lora(self, name, lora_weights, scale=1.0):        self.loaded_loras[name] = {'weights': lora_weights, 'scale': scale}    def get_merged_update(self):        total_update = None        for name, lora in self.loaded_loras.items():            A = lora['weights']['A']            B = lora['weights']['B']            scaling = lora['weights'].get('scaling', 1.0)            update = (A @ B * scaling).T * lora['scale']            if total_update is None:                total_update = update            else:                total_update = total_update + update        return total_update    def list_loras(self):        return [{'name': n, 'scale': l['scale']} for n, l in self.loaded_loras.items()] def apply_lora_to_model(model, target_modules, rank=16, alpha=32):    injected = []    for name, module in model.named_modules():        if any(t in name for t in target_modules) and isinstance(module, nn.Linear):            parent_name = '.'.join(name.split('.')[:-1])            child_name = name.split('.')[-1]            parent = dict(model.named_modules())[parent_name]            lora_layer = LoRALayer(module, rank=rank, alpha=alpha)            setattr(parent, child_name, lora_layer)            injected.append(name)    return injected if __name__ == '__main__':    base_model = nn.Sequential(nn.Linear(768, 768), nn.ReLU(), nn.Linear(768, 768))    injected = apply_lora_to_model(base_model, ['0', '2'], rank=16, alpha=32)    print(f'Injected LoRA to: {injected}')    total_params = sum(p.numel() for p in base_model.parameters())    trainable = sum(p.numel() for p in base_model.parameters() if p.requires_grad)    print(f'Total: {total_params}, Trainable: {trainable}, Ratio: {trainable/total_params*100:.2f}%')    manager = LoRAManager()    manager.register_lora('watercolor', {'A': torch.randn(768,16), 'B': torch.zeros(16,768), 'scaling': 2.0}, scale=0.7)    print(f'LoRAs: {manager.list_loras()}')