torch onnx 权重对比一致性

torch与onnx各层权重layer weight提取
torch与onnx权重layer weight 对比是否一致

def compare_cos_sim(tensor1, tensor2):
	t1 = tensor1.reshape(-1)
  	t2 = tensor2.reshape(-1)
  	n1 = np.linalg.norm(t1)
  	n2 = np.linalg.norm(t2)
  	cos = np.dot(t1/n1, t2.T/n2)
  	print('cos:',cos)
  	


###    torch layer weight
import torch
from models import Net
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
pth_path='epoch-best.pth' 
torch_model = Net()  # 构造模型,创建新模型,网络实例
loaded_model = torch.load(pth_path)  # 加载模型参数
torch_model.load_state_dict(loaded_model['state_dict'])  #将模型参数加载到构造的新建模型实例model_demo中,需要创建的model_demo模型和加载模型的结构、参数名称、参数维度相同,不同时,可选择加载相同部分参数
torch_model = torch_model.to(device)
torch_model.eval()
for name in model.state_dict():                                             
    if name=='wh.2.weight':  
        print(name,':',model.state_dict()[name].shape)
        print('min max:',model.state_dict()[name].min().numpy(),model.state_dict()[name].max().numpy())



###    onnx layer weight
import onnx
from onnx import numpy_helper
onnx_path='epoch-best.onnx'
onnx_model = onnx.load(onnx_path)
for initializer in onnx_model.graph.initializer:
    if initializer.name ==  'wh.2.weight':
        W= numpy_helper.to_array(initializer)
        print(initializer.name,':',W.shape)
        print('min max:',np.min(W),np.max(W))
 
 
       
###    torch与onnx layer weight 相似度计算
keys_list = list(model.state_dict().keys())
print('keys_list',keys_list)

for i, initializer in enumerate(onnx_model.graph.initializer):   
    if initializer.name in keys_list: 
        torch_layer_weight = model.state_dict()[initializer.name]
        onnx_layer_weight= numpy_helper.to_array(initializer)
        compare_cos_sim(torch_layer_weight,onnx_layer_weight)

你可能感兴趣的:(深度学习,pytorch)