YOLOv3自定义2类模型.weights转.pb时报reshape相关ValueError求助

你遇到的ValueError: cannot reshape array of size 278540 into shape (256,128,3,3)报错，核心原因是转换工具默认适配COCO数据集80类的YOLOv3结构，你训练的是2类自定义模型，输出层的参数维度和默认结构不匹配导致权重加载失败，按以下步骤操作即可解决：

1. 核验配置文件参数 #

这是最常见的出错原因，首先确认你训练用的cfg配置文件修改完全正确：

• 找到cfg中全部3个[yolo]层，每个层下的classes参数值都改为2
• 找到每个[yolo]层上方紧邻的[convolutional]层，将每个卷积层的filters参数值改为21，计算公式为3*(类别数 + 5) = 3*(2+5) =21，注意3个yolo层对应的卷积层都要修改，仅修改单一层会直接触发维度不匹配报错
• 转换时使用的cfg文件必须和训练时的版本完全一致，不要混用官方原版YOLOv3 cfg

2. 调整转换工具的类别配置 #

你之前试用的所有转换仓库均默认适配80类，转换前需手动修改相关配置：

• 确认工具加载类别文件的路径指向你自己的2类classes.names，不要用默认的coco.names
• 部分仓库存在硬编码的类别数常量，全局搜索代码中数值为80的类别参数，改为2后再运行转换脚本

3. 标准转换流程参考 #

Darknet权重转TensorFlow SavedModel #

python convert.py --weights ./你的权重文件名.weights --output ./checkpoints/yolov3-custom --num_classes 2 --model yolov3

SavedModel转PB格式 #

python -m tensorflow.python.tools.freeze_graph \
--input_saved_model_dir=./checkpoints/yolov3-custom \
--output_node_names='yolov3/yolo_nms/Reshape_1,yolov3/yolo_nms/Reshape_2' \
--output_graph=./yolov3-custom.pb

PB转TFLite格式 #

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_saved_model('./checkpoints/yolov3-custom')
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]
tflite_model = converter.convert()

with open('yolov3-custom.tflite', 'wb') as f:
  f.write(tflite_model)

异常排查 #

如果修改后仍报错，按以下顺序核验：

• 确认你使用的权重是训练完成后导出的最终权重，不是训练中断生成的临时权重
• 转换前清空所有缓存的checkpoint、中间模型文件，避免旧文件干扰
• 确认你的TensorFlow版本在2.4~2.10区间，过高或过低版本都会出现算子不兼容问题

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