TensorFlow.js网页端MRI预测偶发片段着色器编译失败报错排查

问题说明 #

运行基于TensorFlow.js的MRI影像预测网页时，偶发抛出Uncaught (in promise) Error: Failed to compile fragment shader错误，该错误无固定触发规律，重启Live Server后可能自行恢复，报错中文含义为未捕获的Promise异常：片段着色器编译失败，属于WebGL后端GPU侧运行错误。

故障根因 #

• 模型加载逻辑不合理：当前代码每次用户上传图片都会重新加载模型，重复触发WebGL着色器编译流程，和浏览器的图片解码、DOM渲染进程抢占GPU资源，偶发资源未就绪就执行预测导致编译失败；同时重复加载模型未释放旧张量，会逐步占满WebGL显存，触发异常。
• 图像预处理逻辑错误：代码直接将任意尺寸的上传图片通过tf.reshape强制转为[1,200,200,3]维度，没有做显式的尺寸缩放，当上传图片原始尺寸不是200*200时，张量维度不匹配会直接触发WebGL层异常。
• 开发环境缓存异常：Live Server默认无强缓存校验，偶发出现模型分片、TFJS后端脚本加载不完整的情况，导致着色器初始化资源缺失。

修复方案 #

• 页面初始化阶段预加载模型，模型加载完成后再开启文件上传权限，禁止每次预测重复加载模型
• 修正图像预处理逻辑：上传图片后显式缩放至模型要求的200*200尺寸，完成像素值归一化后再传入模型，预测完成后及时销毁中间张量释放WebGL内存
• 增加WebGL后端初始化容错，异常时自动重试一次

修复后参考代码 #

<script>
let model = undefined;
// 替换为实际的脑肿瘤分类类别列表
const TARGET_CLASSES_BRAIN = [];
// 页面加载完成即预加载模型
window.addEventListener('load', async () => {
  try {
    model = await tf.loadLayersModel("http://127.0.0.1:5500/modelsBrain/modelBrain.json");
    // 预热模型，提前完成着色器编译，避免首次预测卡顿/异常
    model.predict(tf.zeros([1,200,200,3])).dispose();
    document.getElementById("file-input").disabled = false;
  } catch (e) {
    console.error("模型加载失败，重试一次：", e);
    model = await tf.loadLayersModel("http://127.0.0.1:5500/modelsBrain/modelBrain.json");
    document.getElementById("file-input").disabled = false;
  }
})

async function runPrediction(){
  if (!model) return alert("模型未加载完成，请稍候");
  const image = document.getElementById("image");
  // 自动清理中间张量，避免WebGL内存泄漏
  const predResult = tf.tidy(() => {
    const imageTensor = tf.browser.fromPixels(image)
      .toFloat()
      // 按训练时的归一化逻辑调整，若训练时未做0-1归一化可删除下一行
      .div(255.0)
      .resizeNearestNeighbor([200,200])
      .expandDims(0);
    const prediction = model.predict(imageTensor);
    return prediction.softmax().dataSync();
  });

  let topResult = Array.from(predResult)
    .map((p, i) => ({probability: p, className: TARGET_CLASSES_BRAIN[i]}))
    .sort((a,b) => b.probability - a.probability)
    .slice(0,4);
  
  topResult.forEach(p => {
    alert(`${p.className} 预测概率：${(p.probability*100).toFixed(2)}%`);
  })
}

const fileInput = document.getElementById("file-input");
const image = document.getElementById("image");
// 初始禁用上传按钮，模型加载完成后再放开
fileInput.disabled = true;

function getImage() {
  if(!fileInput.files[0]) throw new Error("未选择图片");
  const file = fileInput.files[0];
  const reader = new FileReader();
  reader.onload = function (event) {
    const dataUrl = event.target.result;
    const imageElement = new Image();
    imageElement.src = dataUrl;
    imageElement.onload = async function () {
      image.setAttribute("src", this.src);
      image.setAttribute("height", this.height);
      image.setAttribute("width", this.width);
      await runPrediction();
    };
  };
  reader.readAsDataURL(file);
}
fileInput.addEventListener("change", getImage);
</script>

验证方法 #

如果修改后仍偶发异常，可以在脚本最开头加一行tf.setBackend('cpu')强制使用CPU后端运行预测，若CPU模式下完全无报错，即可确认是WebGL资源抢占/显存泄漏导致的问题，上述修复代码已覆盖对应场景。

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