使用TensorSpace可视化TensorFlow.js预训练模型

原创于 2025-06-08 09:09:40 发布 · 1k 阅读

CC 4.0 BY-SA版权

使用TensorSpace可视化TensorFlow.js预训练模型

tensorspace Neural network 3D visualization framework, build interactive and intuitive model in browsers, support pre-trained deep learning models from TensorFlow, Keras, TensorFlow.js 项目地址: https://gitcode.com/gh_mirrors/te/tensorspace

TensorSpace是一个用于神经网络可视化的开源库，它能够将复杂的深度学习模型以3D交互形式呈现。本文将详细介绍如何利用TensorSpace和配套的TensorSpace-Converter工具来可视化TensorFlow.js预训练模型。

准备工作

在开始之前，我们需要准备以下内容：

一个训练好的TensorFlow.js模型（通常包含一个描述模型结构的.json文件和一个存储权重的.bin文件）
安装好的TensorSpace和TensorSpace-Converter环境
了解模型的基本结构，特别是各层的名称

本文以经典的LeNet模型（在MNIST数据集上训练）为例进行说明。如果你还没有现成的模型，可以使用TensorFlow.js自行训练一个简单的LeNet模型。

模型转换流程

TensorSpace-Converter是将原始模型转换为TensorSpace可识别格式的关键工具。对于TensorFlow.js模型，转换步骤如下：

准备模型文件：确保模型的.json和.bin文件位于同一目录下
确定输出层名：需要提取模型中各层的名称，这些名称将用于后续可视化
执行转换命令：

tensorspace_converter \
    --input_model_from="tfjs" \
    --output_layer_names="myPadding,myConv1,myMaxPooling1,myConv2,myMaxPooling2,myDense1,myDense2,myDense3" \
    ./rawModel/mnist.json \
    ./convertedModel/

关键参数说明：

input_model_from：指定输入模型来源为"tfjs"
output_layer_names：列出模型中所有需要可视化的层名称
最后一个参数指定输出目录

转换完成后，会在目标目录生成两个文件：

.json：转换后的模型结构描述文件
.bin：转换后的权重文件

可视化实现

转换完成后，我们可以使用TensorSpace API来构建可视化模型。以下是核心代码示例：

// 创建序列模型容器
let model = new TSP.models.Sequential(modelContainer);

// 逐层添加模型结构
model.add(new TSP.layers.GreyscaleInput());  // 灰度输入层
model.add(new TSP.layers.Padding2d());      // 填充层
model.add(new TSP.layers.Conv2d());         // 卷积层1
model.add(new TSP.layers.Pooling2d());      // 池化层1
model.add(new TSP.layers.Conv2d());         // 卷积层2
model.add(new TSP.layers.Pooling2d());      // 池化层2
model.add(new TSP.layers.Dense());          // 全连接层1
model.add(new TSP.layers.Dense());          // 全连接层2
model.add(new TSP.layers.Output1d({         // 输出层
    outputs: ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]
}));

// 加载转换后的模型
model.load({
    type: "tfjs",
    url: "./convertedModel/model.json"
});

// 初始化可视化模型
model.init();