LeakyRelu是否有TensorFlow.js实现？转换后YOLO模型加载报错如何解决

报错原因 #

你在转换模型时使用了tfjs_graph_model输出格式，该格式会将Keras模型的所有层逻辑直接固化为静态计算图节点，不会做算子适配转换。当前你使用的TensorFlow.js版本的WebGL后端未内置leakyrelu算子的实现，因此加载模型时遇到该算子节点无法执行就会抛出报错。
你测试CPU后端也失败大概率是因为切换后端的时机错误：必须在调用loadGraphModel之前完成后端设置并等待设置生效，否则模型加载时还是会调用默认的WebGL后端执行。

WebGL后端未实现LeakyRelu的含义 #

TensorFlow.js的WebGL后端依赖GLSL着色器程序调度GPU完成张量计算，每一个算子都需要有对应的着色器代码实现才能在WebGL环境下运行。所谓未实现就是当前版本TFJS的WebGL后端代码库中，没有为leakyrelu算子编写对应的GPU执行逻辑，因此无法调度GPU执行该算子的计算。

无需重新训练模型的修复方案 #

• 方案1：正确切换CPU后端运行
  升级到最新稳定版TFJS，在加载模型前先完成后端切换，示例代码如下：

  import * as tf from '@tensorflow/tfjs';
  // 先设置后端，等待设置完成后再加载模型
  await tf.setBackend('cpu');
  // 确认后端切换成功
  if (tf.getBackend() === 'cpu') {
    const model = await tf.loadGraphModel("http://127.0.0.1:8080/model.json");
  }
  

  最新版TFJS的CPU后端已完整支持leakyrelu算子，不需要修改模型即可正常运行。

• 方案2：修改Keras模型的激活实现后重新转换
  LeakyRelu的计算公式为max(x, alpha * x)，可以在不重训的前提下将模型中所有的LeakyRelu激活替换为基础算子组合，再重新转换模型即可，示例Keras侧代码：

  from tensorflow.keras.layers import Lambda
  import tensorflow.keras.backend as K
  
  # 定义等效的LeakyRelu实现
  def custom_leaky_relu(x, alpha=0.1):
      return K.maximum(x, alpha * x)
  
  # 遍历替换模型中所有LeakyRelu激活
  for layer in model.layers:
      if hasattr(layer, 'activation') and layer.activation.__name__ == 'leaky_relu':
          layer.activation = Lambda(custom_leaky_relu)
  
  # 保存修改后的模型，再用原来的转换命令转成TFJS格式即可
  model.save('modified_yolo.h5')
  

  修改后转换出的模型计算图中不存在leakyrelu算子节点，WebGL和CPU后端都可以直接运行。

• 方案3：TFJS侧自定义WebGL算子实现
  如果需要保留WebGL的GPU加速能力，可以自行编写LeakyRelu对应的GLSL着色器逻辑，在TFJS中注册自定义算子，加载模型时会自动调用自定义的算子实现，不需要修改原模型。

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