TensorFlow.js中如何复制模型训练权重及转换模型为字符串？

1. 模型权重复制最优方案 #

针对同结构模型的权重复制，官方原生的getWeights()/setWeights()是性能最优的方案，没有额外的序列化开销：

• 全量权重复制（要求两个模型层结构、张量维度完全一致，否则会抛出维度不匹配错误）

// 从源模型读取权重
const sourceWeights = sourceModel.getWeights();
// 直接赋值到目标模型
targetModel.setWeights(sourceWeights);

• 按层匹配复制（适配两个模型局部结构一致的场景）

sourceModel.layers.forEach((sourceLayer, index) => {
  // 优先按层名匹配，无匹配项则按索引匹配
  const targetLayer = targetModel.getLayer(sourceLayer.name, index);
  if (targetLayer) {
    targetLayer.setWeights(sourceLayer.getWeights());
  }
});

2. 模型/权重序列化转字符串（跨Worker传输用） #

跨Worker传输分两种场景选择对应方案：

场景1：仅传输权重（性能最优，目标Worker已预定义相同结构的模型） #

使用TF.js内置的tf.io.encodeWeights/tf.io.decodeWeights接口完成权重的序列化与反序列化，需要字符串格式时可转为base64：

// 源Worker侧：序列化权重
const sourceWeights = sourceModel.getWeights();
const { data: weightBuffer, specs: weightSpecs } = await tf.io.encodeWeights(sourceWeights);
// 转base64字符串
const weightBase64 = btoa(String.fromCharCode(...new Uint8Array(weightBuffer)));
const specsStr = JSON.stringify(weightSpecs);
// 发送到目标Worker
postMessage({ weightBase64, specsStr });

// 目标Worker侧：反序列化并赋值
const weightUint8 = Uint8Array.from(atob(weightBase64), c => c.charCodeAt(0));
const decodedWeights = await tf.io.decodeWeights(weightUint8.buffer, JSON.parse(specsStr));
targetModel.setWeights(decodedWeights);

场景2：传输完整模型（结构+权重，目标Worker无预定义模型） #

直接序列化完整模型 artifacts 为JSON字符串：

// 源Worker侧：序列化完整模型
const modelArtifacts = await sourceModel.save(tf.io.withSaveHandler(async artifacts => artifacts));
const modelStr = JSON.stringify(modelArtifacts);
postMessage({ modelStr });

// 目标Worker侧：还原模型
const loadedArtifacts = JSON.parse(modelStr);
const targetModel = await tf.loadLayersModel(tf.io.fromMemory(loadedArtifacts));

注意事项 #

• 序列化前可调用await tf.ready()确保所有权重计算完成，避免拿到未就绪的张量
• Web Worker的postMessage本身支持结构化克隆ArrayBuffer，若不需要强制字符串格式，直接传输二进制buffer性能更高

内容的提问来源于stack exchange，提问作者Игорь Токарев