遇到一个问题是自己创建了新的Detect检测头,但是在导出模型时,想要修改输出格式,在yolo中可以通过if self.export:
来修改网络的返回值格式
当使用model.export()导出时,理论上会自动将export
设置为True
但是在实际中发现export=false
,于是通过调试发现在ultralytics/engine/exporter.py
中
for m in model.modules():
if isinstance(m, (Detect, RTDETRDecoder)): # includes all Detect subclasses like Segment, Pose, OBB
m.dynamic = self.args.dynamic
m.export = True
m.format = self.args.format
elif isinstance(m, C2f) and not any((saved_model, pb, tflite, edgetpu, tfjs)):
# EdgeTPU does not support FlexSplitV while split provides cleaner ONNX graph
m.forward = m.forward_split
需要对自己的检测头进行实例,加入到代码中,如下
if isinstance(m, (Detect, RTDETRDecoder,MyDetect)):
再次执行model.export()
self.export则为True
但是这个方法并不是最规范的,设置自己的类继承Detect会更好MyDetect(Detect)
参考地址:github