ONNX Runtime Web中WebGPU GPU内存无法释放问题求助

问题根源 #

虽然你已经调用了session.release()和tensor.dispose()，但GPU内存未释放通常是因为Session引用未被彻底清除、Tensor资源未完全释放，或者WebGPU的内存回收机制需要配合JS垃圾回收触发。

具体修正步骤 #

• 确保Session资源彻底释放：调用session.release()后，必须解除对Session实例的引用，让JavaScript垃圾回收器能回收相关内存。
• 提前拷贝Tensor数据：在释放Session前，先将Tensor的CPU数据拷贝出来，避免Session释放后Tensor的GPU关联资源无法正常回收。
• 检查Session创建配置：如果创建Session时启用了内存池（如memArenaOptions），可能会导致内存滞留，可尝试禁用内存池。
• 触发JS垃圾回收：在所有资源释放后，可手动触发垃圾回收（仅调试环境可用），帮助浏览器回收闲置内存。

修改后的代码示例 #

修正infer2函数

/**
 * @param model 模型路径，每次推理创建新Session，避免内存泄漏
 * @param inputTensor 输入张量
 */
export async function infer2(model: string, inputTensor: Tensor) {
  let session = await newSession(model);
  try {
    const feeds: any = {};
    const inputNames = session.inputNames;
    feeds[inputNames[0]] = inputTensor;
    const results = await session.run(feeds);
    const outputTensor = results[session.outputNames[0]];
    
    // 提前拷贝输出数据到CPU，避免Session释放后数据丢失
    const outputData = Array.from(outputTensor.data as Float32Array);
    const outputDims = [...outputTensor.dims];
    
    // 立即释放输出Tensor的GPU资源
    outputTensor.dispose();
    
    return { data: outputData, dims: outputDims };
  } finally {
    // 确保Session释放完成
    await session.release();
    // 解除Session引用，帮助GC回收
    session = null;
  }
}

修正infer函数

/**
 * 加载ONNX模型并执行推理
 * @param model 模型路径
 * @param input 输入Ndarray
 * @returns 输出Ndarray
 */
export const infer = async (model: string, input: Ndarray) => {
  const inputTensor = ndarrayToTensor(input);
  try {
    const { data, dims } = await infer2(model, inputTensor);
    return new Ndarray(data as number[], dims as number[]);
  } finally {
    // 释放输入Tensor的GPU资源
    inputTensor.dispose();
  }
};

额外注意事项 #

• nvidia-smi显示的GPU内存占用可能包含浏览器进程的整体占用，而非仅ONNX Runtime的内存。可关闭其他标签页后重新测试，或使用浏览器的DevTools（如Chrome的Memory面板）查看JS内存和GPU内存使用情况。
• WebGPU的内存回收由浏览器调度，即使代码正确释放资源，内存可能不会立即下降，需等待浏览器的垃圾回收周期触发。

内容的提问来源于stack exchange，提问作者chikadance