PyTorch神经网络打印存储所有权重+激活值(运行时中间值)

很多时候嵌入式或者新硬件需要纯净的权重模型和激活值(运行时中间值),本文提供一种最简洁的方法。

假设已经有模型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作为一个整体,目前没办法拆开来看其内部的中间值。

相关推荐
颜颜yan_3 分钟前
Python面向对象编程详解:从零开始掌握类的声明与使用
开发语言·redis·python
闻道且行之25 分钟前
Windows|CUDA和cuDNN下载和安装,默认安装在C盘和不安装在C盘的两种方法
windows·深度学习·cuda·cudnn
jonyleek1 小时前
如何搭建一套安全的,企业级本地AI专属知识库系统?从安装系统到构建知识体系,全流程!
人工智能·安全
我的ID配享太庙呀1 小时前
Django 科普介绍:从入门到了解其核心魅力
数据库·后端·python·mysql·django·sqlite
MQ_SOFTWARE2 小时前
AI驱动的金融推理:Fin-R1模型如何重塑行业决策逻辑
人工智能·金融
@蓝莓果粒茶2 小时前
LeetCode第350题_两个数组的交集II
c++·python·学习·算法·leetcode·职场和发展·c#
生医转码,四海为家2 小时前
零基础-动手学深度学习-6.6 卷积神经网络(LeNet)
人工智能·深度学习·cnn
无名工程师2 小时前
AI 学习过程中各阶段的学习重点、时间规划以及不同方向的选择与建议等内容
人工智能·学习
FinAnalyzer2 小时前
如何在 InsCodeAI 上搭建并使用 Jupyter Notebook 环境?
ide·python·jupyter
java1234_小锋2 小时前
【NLP舆情分析】基于python微博舆情分析可视化系统(flask+pandas+echarts) 视频教程 - 微博文章数据可视化分析-文章分类下拉框实现
python·自然语言处理·flask