diff --git a/test/resources/model.bin b/test/resources/model.bin new file mode 100644 index 0000000..d796950 Binary files /dev/null and b/test/resources/model.bin differ diff --git a/tools/export_model_weights.py b/tools/export_model_weights.py new file mode 100644 index 0000000..e9f839e --- /dev/null +++ b/tools/export_model_weights.py @@ -0,0 +1,28 @@ +import torch +import struct + +def export_model_weights(model: torch.nn.Module, filename): + with open(filename, 'wb') as f: + + header = "" + offset = 0 + + + for name, param in model.named_parameters(): + if 'weight' not in name and 'bias' not in name: + continue + + tensor_values = param.flatten().tolist() + tensor_bytes = struct.pack('f' * len(tensor_values), *tensor_values) + + tensor_size = param.numel() + + header += f"{name},{tensor_size},{offset}\n" + + offset += len(tensor_bytes) + + f.write(tensor_bytes) + + f.seek(0) + f.write(struct.pack('q', len(header))) + f.write(header.encode('utf-8')) \ No newline at end of file diff --git a/tools/model_test.py b/tools/model_test.py index bafb57e..848153f 100644 --- a/tools/model_test.py +++ b/tools/model_test.py @@ -1,6 +1,8 @@ import torch import utils +from export_model_weights import export_model_weights + class TestModel(torch.nn.Module): def __init__(self, *args, **kwargs) -> None: @@ -121,3 +123,5 @@ if __name__ == "__main__": ]).reshape(2, 6, 6) out = model(input) utils.print_cpp_vector(out) + + export_model_weights(model, "test/resources/model.bin")