目的是得到以下模型1、官方yolox_s的源码和yolox_s.pth获取https://github.com/Megvii-BaseDetection/YOLOX2、修改yolo_head.py的forward替换为以下defforward(self,xin,labelsNone,imgsNone):outputs[]fork,(cls_conv,reg_conv,stride_this_level,x)inenumerate(zip(self.cls_convs,self.reg_convs,self.strides,xin)):xself.stems[k](x)cls_featcls_conv(x)reg_featreg_conv(x)cls_outputself.cls_preds[k](cls_feat)# [B, C, H, W]reg_outputself.reg_preds[k](reg_feat)# [B, 4, H, W]obj_outputself.obj_preds[k](reg_feat)# [B, 1, H, W]# 关键不要 decode不要 concatoutputs.append(reg_output)outputs.append(obj_output)outputs.append(cls_output)returnoutputs3、修改export_onnx.py的main()为以下defmain():argsmake_parser().parse_args()logger.info(args value: {}.format(args))expget_exp(args.exp_file,args.name)exp.merge(args.opts)ifnotargs.experiment_name:args.experiment_nameexp.exp_name modelexp.get_model()ifargs.ckptisNone:file_nameos.path.join(exp.output_dir,args.experiment_name)ckpt_fileos.path.join(file_name,best_ckpt.pth)else:ckpt_fileargs.ckpt# load the model state dictckpttorch.load(ckpt_file,map_locationcpu)model.eval()ifmodelinckpt:ckptckpt[model]model.load_state_dict(ckpt)modelreplace_module(model,nn.SiLU,SiLU)model.head.decode_in_inferenceFalselogger.info(loading checkpoint done.)dummy_inputtorch.randn(args.batch_size,3,exp.test_size[0],exp.test_size[1])output_names[]output_names[reg1,obj1,cls1,reg2,obj2,cls2,reg3,obj3,cls3,]torch.onnx._export(model,dummy_input,args.output_name,input_names[args.input],output_namesoutput_names,dynamic_axes{args.input:{0:batch},**{name:{0:batch}fornameinoutput_names}}ifargs.dynamicelseNone,opset_versionargs.opset,)logger.info(generated onnx model named {}.format(args.output_name))ifnotargs.no_onnxsim:importonnxfromonnxsimimportsimplify# use onnx-simplifier to reduce reduent model.onnx_modelonnx.load(args.output_name)model_simp,checksimplify(onnx_model)assertcheck,Simplified ONNX model could not be validatedonnx.save(model_simp,args.output_name)logger.info(generated simplified onnx model named {}.format(args.output_name))4、导出指令python tools/export_onnx.py-fexps/default/yolox_s.py-cyolox_s.pth --output-name yolox_s.onnx--opset12--output.上述完成就可得到需要的onnx