webassembly003 MINISIT mnist/convert-h5-to-ggml.py + ggml format

数据结构

# Convert MNIS h5 transformer model to ggml format
#
# Load the (state_dict) saved model using PyTorch
# Iterate over all variables and write them to a binary file.
#
# For each variable, write the following:
#   - Number of dimensions (int)
#   - Name length (int)
#   - Dimensions (int[n_dims])
#   - Name (char[name_length])
#   - Data (float[n_dims])
#
# At the start of the ggml file we write the model parameters

• 这个简单的版本没有Name的部分，导出的数据最终如下

ggml-model-f32.bin	注释
0x67676d6c	magic number (filesignatures)
2	len(fc1.weight.shape)
784	fc1.weight.shape = (500, 784)
500	fc1.weight.shape = (500, 784)
data	fc1.weight
1	len(fc1.bias.shape)
500	fc1.bias.shape = (500, )
data	fc1.bias
2	len(fc2.weight.shape)
500	fc1.weight.shape = (10, 500)
10	fc1.weight.shape =(10, 500)
data	fc2.weight
1	len(fc2.bias.shape)
10	fc2.bias.shape =(10,)
data	fc1.bias

代码注释

import sys
import struct
import json
import numpy as np
import re

import torch
import torch.nn as nn
import torchvision.datasets as dsets
import torchvision.transforms as transforms
from torch.autograd import Variable

# 检查是否提供了正确数量的命令行参数
if len(sys.argv) != 2:
    print("Usage: convert-h5-to-ggml.py model\n")
    sys.exit(1)

# 获取输入h5模型和输出ggml模型的文件路径
state_dict_file = sys.argv[1]
fname_out = "models/mnist/ggml-model-f32.bin"

# 加载PyTorch保存的state_dict模型
state_dict = torch.load(state_dict_file, map_location=torch.device('cpu'))

# 以写入模式打开输出二进制文件
fout = open(fname_out, "wb")

# 在文件中写入魔术数字'ggml'，以十六进制格式作为文件标识符
# 使用 Python 的 struct 模块将整数 0x67676d6c 打包为二进制数据的操作。在这里，"i" 表示使用整数格式进行打包。
fout.write(struct.pack("i", 0x67676d6c))  # magic: ggml in hex 

# 迭代state_dict中的所有变量
for name in state_dict.keys():
    # 从变量中提取数据并将其转换为NumPy数组
    data = state_dict[name].squeeze().numpy()
    print("Processing variable: " + name + " with shape: ", data.shape) 
    n_dims = len(data.shape);
   
    # 将变量的维度数量写入二进制文件
    fout.write(struct.pack("i", n_dims))
    
    # 将数据转换为float32并将维度写入二进制文件
    data = data.astype(np.float32)
    for i in range(n_dims):
        fout.write(struct.pack("i", data.shape[n_dims - 1 - i]))

    # 将数据写入二进制文件
    data.tofile(fout)

# 关闭二进制文件
fout.close()

print("Done. Output file: " + fname_out)
print("")

tofile()

• NumPy提供的存数组内容的文件操作函数。读取使用fromfile。

struct.pack

• 将字节解释为打包的二进制数据。

输出

$:~/ggml/ggml/examples/mnist$ python3 ./convert-h5-to-ggml.py 
./models/mnist/mnist_model.state_dictOrderedDict([('fc1.weight', tensor([[ 0.0130,  0.0034, -0.0287,  ..., -0.0268, -0.0352, -0.0056],
        [-0.0134,  0.0077, -0.0028,  ...,  0.0356,  0.0143, -0.0107],
        [-0.0329,  0.0154, -0.0167,  ...,  0.0155,  0.0127, -0.0309],
        ...,
        [-0.0216, -0.0302,  0.0085,  ...,  0.0301,  0.0073,  0.0153],
        [ 0.0289,  0.0181,  0.0326,  ...,  0.0107, -0.0314, -0.0349],
        [ 0.0273,  0.0127