通义万相2.1的LoRA怎么训练-从原理到实战指南

W' = W + ΔW = W + BA^T 

# 创建Python虚拟环境
conda create -n wanxiang-lora python=3.10
conda activate wanxiang-lora

# 安装关键库
pip install torch==2.1.0+cu121 -f https://download.pytorch.org/whl/torch_stable.html
pip install diffusers transformers accelerate peft xformers
pip install datasets pillow tensorboard

from diffusers import StableDiffusionPipeline

model_path = "wanxiang/wanxiang-v2.1"
pipe = StableDiffusionPipeline.from_pretrained(model_path)

{subject} {action}, {art_style} style, 
{lighting}, {composition}, 
detailed {texture}, color scheme: {colors}

from albumentations import *
transform = Compose([  RandomResizedCrop(512, 512, scale=(0.8, 1.0)),  HorizontalFlip(p=0.5),  ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),  GaussNoise(var_limit=(10, 50)),
])

# lora_config.yaml
rank: 64              # 核心维度参数 (8-128)
alpha: 32             # 缩放因子 (通常=rank)
target_modules:        # 注入位置- "to_k"- "to_v"- "to_q"- "ff.net.0.proj"
dropout: 0.05
bias: "none"

from peft import LoraConfig, get_peft_model

# 创建LoRA配置
lora_config = LoraConfig(  r=args.rank,  lora_alpha=args.alpha,  target_modules=target_modules,  lora_dropout=args.dropout
)

# 注入LoRA到模型
model.unet = get_peft_model(model.unet, lora_config)

# 优化器配置
optimizer = torch.optim.AdamW(  model.unet.parameters(),  lr=1e-4,  weight_decay=1e-4
)

# 训练循环
for epoch in range(epochs):  for batch in dataloader:      clean_images = batch["images"]      latents = vae.encode(clean_images).latent_dist.sample()      noise = torch.randn_like(latents)      timesteps = torch.randint(0, 1000, (len(latents),))            noisy_latents = scheduler.add_noise(latents, noise, timesteps)      noise_pred = model.unet(noisy_latents, timesteps).sample            loss = F.mse_loss(noise_pred, noise)      loss.backward()      optimizer.step()      optimizer.zero_grad()

from diffusers import StableDiffusionPipeline
import torch

pipeline = StableDiffusionPipeline.from_pretrained(  "wanxiang/wanxiang-v2.1",  torch_dtype=torch.float16
)
pipeline.unet.load_attn_procs("lora_weights.safetensors")
pipeline.to("cuda")

# 生成图像
image = pipeline(  "A robot painting in Van Gogh style, lora_weight=0.8",  guidance_scale=7.5,  num_inference_steps=50
).images[0]

# 将LoRA权重合并到基础模型
merged_model = pipeline.unet
for name, module in merged_model.named_modules():  if hasattr(module, "merge_weights"):      module.merge_weights(merge_alpha=0.85)  # 融合比例调节

# 保存完整模型
merged_model.save_pretrained("wanxiang_van_gogh_robot")

# 启用8-bit优化器
accelerate launch --config_file config.yaml train.py \--use_8bit_adam

# 梯度累积技术
training_args = TrainingArguments(  per_device_train_batch_size=1,  gradient_accumulation_steps=4,
)

# 混合精度训练
torch.cuda.amp.autocast(enabled=True)

rank: 96
steps: 2000
lr_scheduler: cosine_with_warmup
lr_warmup_steps: 100
prompt_template: "best quality, masterpiece, illustration, [character_name]"

原始模型： "an anime girl with blue hair"

+ LoRA后："Skye from Neon Genesis, aqua hair with glowing tips,  mecha suit design, signature pose, studio Ghibli background"

# 监控GPU状态
watch -n 1 nvidia-smi

# 启动TensorBoard
tensorboard --logdir=./logs --port 6006