一种将torch模型权重迁移到C++中并计算的方法

开场白

本文介绍一种将torch模型权重迁移到C++中并计算的方法。项目代码见链接。如有错误或需改善，请别吝啬您的意见。

numpy转为二进制类型，tensor转为ndarray

torch可以将tensor转为ndarray。numpy可以将array转为二进制数据。

>>> import torch
>>> import numpy as np
>>> import strct
>>> x = torch.randn(3, 4).numpy().astype(np.float32)
>>> struct.unpack("12f", x.tobytes())

torch权重字典为ordereddict

torch权重字典为ordereddict。

>>> type(torch.nn.Module().state_dict())
<class 'collections.OrderedDict'>

部分提取键与值并写入的流程

这是本视频将权重写为二进制数据的基本思路。将一对键值以尖括号括起来。

配置文件与二进制文件

这个配置文件中，值是键的长度。

C++中读取、显示float32buffer

定义好接收的缓冲区，确保数据对应上，这里采用结构体。分别读取地址写入数据。




可以参考上方代码简单输出查看下。

矩阵计算并输出，与numpy对比

为了不依赖额外第三方库，简单实现了一套矩阵乘法计算库。由于精度与算法，结果有一定误差。

读取bert权重并导出二进制文件与配置文件

这是模型提取的一个例子。





上四图分别为头尾部分的配置文件和二进制文件。

总结

原理比较简单，本质就是有序字典数据转为二进制数据。