本文以bert为例,对比了添加Lora模块前后的网络结构图
说明:

• 1.为了加快速度,将bert修改为一层
• 2.lora只加到intermediate.dense,方便对比
• 3.使用了几种不同的可视化方式(onnx可视化,torchviz图,torch.fx可视化,tensorboard可视化)

可参考的点:

• 1.peft使用
• 2.几种不同的pytorch模型可视化方法

一.效果图

1.torch.fx可视化

A.添加前

B.添加后

2.onnx可视化

A.添加前

B.添加后

3.tensorboard可视化

A.添加前

B.添加后

二.复现步骤

1.生成配置文件(num_hidden_layers=1)

tee ./config.json <<-'EOF'
{"architectures": ["BertForMaskedLM"],"attention_probs_dropout_prob": 0.1,"directionality": "bidi","hidden_act": "gelu","hidden_dropout_prob": 0.1,"hidden_size": 768,"initializer_range": 0.02,"intermediate_size": 3072,"layer_norm_eps": 1e-12,"max_position_embeddings": 512,"model_type": "bert","num_attention_heads": 12,"num_hidden_layers": 1,"pad_token_id": 0,"pooler_fc_size": 768,"pooler_num_attention_heads": 12,"pooler_num_fc_layers": 3,"pooler_size_per_head": 128,"pooler_type": "first_token_transform","type_vocab_size": 2,"vocab_size": 21128
}
EOF

2.运行测试脚本

tee bert_lora.py <<-'EOF'
import time
import os
import torch
import torchvision.models as models
import torch.nn as nn
import torch.nn.init as init
import time
import numpy as np
from peft import get_peft_config, get_peft_model, get_peft_model_state_dict, LoraConfig, TaskType
from torchviz import make_dot
from torch.utils.tensorboard import SummaryWriter
from torch._functorch.partitioners import draw_graphdef onnx_infer_shape(onnx_path):import onnxonnx_model  = onnx.load_model(onnx_path)new_onnx= onnx.shape_inference.infer_shapes(onnx_model)onnx.save_model(new_onnx, onnx_path)def get_model():torch.manual_seed(1)from transformers import AutoModelForMaskedLM,BertConfigconfig=BertConfig.from_pretrained("./config.json")model = AutoModelForMaskedLM.from_config(config)return model,configdef my_compiler(fx_module: torch.fx.GraphModule, _):draw_graph(fx_module, f"bert.{time.time()}.svg")return fx_module.forwardif __name__ == "__main__":model,config=get_model()model.eval()input_tokens=torch.randint(0,config.vocab_size,(1,128))# 一.原始模型# 1.onnx可视化torch.onnx.export(model,input_tokens,"bert_base.onnx",export_params=False,opset_version=11,do_constant_folding=True)onnx_infer_shape("bert_base.onnx")# 2.torchviz图output = model(input_tokens)logits = output.logitsviz = make_dot(logits, params=dict(model.named_parameters()))viz.render("bert_base", view=False)# 3.torch.fx可视化compiled_model = torch.compile(model, backend=my_compiler)output = compiled_model(input_tokens)# 4.tensorboard可视化writer = SummaryWriter('./runs')writer.add_graph(model, input_to_model = input_tokens,use_strict_trace=False)writer.close()# 二.Lora模型peft_config = LoraConfig(task_type=TaskType.CAUSAL_LM,inference_mode=True,r=8,lora_alpha=32,target_modules=['intermediate.dense'],lora_dropout=0.1,)lora_model = get_peft_model(model, peft_config)lora_model.eval()torch.onnx.export(lora_model,input_tokens,"bert_base_lora_inference_mode.onnx",export_params=False,opset_version=11,do_constant_folding=True)onnx_infer_shape("bert_base_lora_inference_mode.onnx")compiled_model = torch.compile(lora_model, backend=my_compiler)output = compiled_model(input_tokens)writer = SummaryWriter('./runs_lora')writer.add_graph(lora_model, input_to_model = input_tokens,use_strict_trace=False)writer.close()
EOF# 安装依赖
apt install graphviz -y
pip install torchviz
pip install pydot# 运行测试程序
python bert_lora.py