很多时候嵌入式或者新硬件需要纯净的权重模型和激活值(运行时中间值),本文提供一种最简洁的方法。
假设已经有模型model和pt文件了,在当前目录下新建weights文件夹,运行这段代码,就可以得到模型的权重(文本形式和二进制形式)
python
model.load_state_dict(state_dict)
global_index = 0
for name, param in model.named_parameters():
print(name, param.size())
print(param.data.numpy(),file=open(f"weights/{global_index}-{name}.txt", "w"))
param.data.numpy().tofile(f"weights/{global_index}-{name}.bin")
global_index += 1
对于二进制形式的文件,可以通过od -t f4 <binary file name>
查看其对应的浮点数值。f4
表示fp32.
打印forward的中间值:(这么复杂是必要的)
python3
global_index = 0
def hook_fn(module, input, output):
global global_index
module_name = str(module)
module_name=module_name.replace(" ", "")
module_name=module_name.replace("\n", "")
# print(name)
intermediate_outputs = {}
# input is a tuple, output is a tensor
for i, inp in enumerate(input):
intermediate_outputs[f"{global_index}-{module_name}-input-{i}"] = inp
intermediate_outputs[f"{global_index}-{module_name}-output"] = output
module_name = module_name[0:200] # make sure full path <= 255
print(intermediate_outputs)
print(f"Size input:",end=" ")
if(type(input) == tuple):
for i, inp in enumerate(input):
if type(inp) == torch.Tensor:
print(f"{i}-th Size: {inp.size()}", end=", ")
inp.numpy().tofile(f"activations/{global_index}-{module_name}-input-{i}.bin")
else:
print(f"{i}-th : {inp}", end=", ")
elif type(input) == torch.Tensor:
print(f"Size: {input.size()}")
input.numpy().tofile(f"activations/{global_index}-{module_name}-input.bin")
print(f"Size output: {output.size()}")
global_index += 1
output.numpy().tofile(f"activations/{global_index}-{module_name}-output.bin")
def register_hooks(model):
for name, layer in model.named_children():
# print(name, layer) # dump all layers, > layers.txt
# Register the hook to the current layer
layer.register_forward_hook(hook_fn)
# Recursively apply the same to all submodules
register_hooks(layer)
register_hooks(model)
其中regster_hooks
和以下等价(不需要recursive了)
python3
def register_hooks(model):
for name, layer in model.named_modules():
# print(name, layer) # dump all layers
layer.register_forward_hook(hook_fn)
其中nn.sequential
作为一个整体,目前没办法拆开来看其内部的中间值。