yolov8的P5特征具有很强的应用的价值之前也分析过使用P5特征做特征比对本文章是实操性的用来将P5特征从yolov8中提取出来这里的“提取来”指的是将embeding的相关信息在yolo-tensorrt中提取。注意有两种情况一种是P5参训练一种是P5不参与训练P5不参与训练完全可以正常训练yolo在export的时候在去修改forward这个时候好像是只需要修改forward中的_predict_once函数就行其他不需要修改可以运行下试试原来forward是一个返回值现在变成了两个有些地方不适配所以报错跟着错误修改就行P5参与训练多任务需求在训练前修改forwardpredict、validator层都会受到影响损失层也需要修改我下面的修改比较适配P5不参与训练P5参与训练就比较复杂了但是也可以改。一、forward 中提取P5特征ultralytics/nn/task.py 中的 _predict_once() 是 forward的核心操作def_predict_once(self,x,profileFalse,visualizeFalse,embedNone):Perform a forward pass through the network. Args: x (torch.Tensor): The input tensor to the model. profile (bool): Print the computation time of each layer if True. visualize (bool): Save the feature maps of the model if True. embed (list, optional): A list of layer indices to return embeddings from. Returns: (torch.Tensor): The last output of the model. y,dt,embeddings[],[],[]# outputsembedfrozenset(embed)ifembedisnotNoneelse{-1}max_idxmax(embed)forminself.model:ifm.f!-1:# if not from previous layerxy[m.f]ifisinstance(m.f,int)else[xifj-1elsey[j]forjinm.f]# from earlier layersifprofile:self._profile_one_layer(m,x,dt)xm(x)# runy.append(xifm.iinself.saveelseNone)# save outputifvisualize:feature_visualization(x,m.type,m.i,save_dirvisualize)ifm.iinembed:embeddings.append(torch.nn.functional.adaptive_avg_pool2d(x,(1,1)).squeeze(-1).squeeze(-1))# flattenifm.imax_idx:returntorch.unbind(torch.cat(embeddings,1),dim0)# 新添加feat21torch.nn.functional.adaptive_avg_pool2d(y[21],(1,1)).squeeze(-1).squeeze(-1)#feat21 F.normalize(feat21, dim1)(最好不要在这里归一化因为转trt权重没有问题但是在bool context m_context-executeV2((void**)bindings)报错了有些tensorrt核好像不适配)returnx,feat21# 修改由return x 修改为 return x, feat21二、接口修正上面的修改不是为了修改predict 而是为了export 的时候多一个暴漏在外的结点这样可以在yolo-tensorrt中通过bidding而获得由于_predict_once函数原来 是return x 修改为 return x, feat21为了适应原代码直接按下面的方式修正接口就行修改DetectionModel 文件ultralytics/nn/tasks.py 找到classDetectionModel(BaseModel):在 __init__ m.stridetorch.tensor([s/x.shape[-2]forxin_forward(torch.zeros(1,ch,s,s))])修改为 m.stridetorch.tensor([s/x.shape[-2]forxin_forward(torch.zeros(1,ch,s,s))[0])修改validator ultralytics/engine/validator.py 中def__call__(self)函数中的第221行 predsself.postprocess(preds)修改为 predsself.postprocess(preds[0])修改predictor ultralytics/engine/predictor.py 中defstream_inference(self)函数中的第335行 self.resultsself.postprocess(preds,im,im0s)修改为 self.resultsself.postprocess(preds[0],im,im0s)三、loss修正第一种修正_predict_once()中多出来的结点不参与训练文件ultralytics/utils/loss.py中找到classv8DetectionLoss的def__call__函数 binary_cls_predNoneif(isinstance(preds,tuple)andlen(preds)2andisinstance(preds[1],torch.Tensor)):preds,binary_cls_predpreds第二种修正_predict_once()中多出来的结点参与训练文件ultralytics/utils/loss.py中找到classv8DetectionLoss的def__call__函数 binary_cls_predNoneif(isinstance(preds,tuple)andlen(preds)2andisinstance(preds[1],torch.Tensor)):preds,binary_cls_predpreds branch_lossFalseifbranch_loss:losstorch.zeros(4,deviceself.device)else:losstorch.zeros(3,deviceself.device)ifbranch_lossandbinary_cls_predisnotNone:binary_targetbatch[has_locomotive].float().to(self.device)binary_lossF.binary_cross_entropy_with_logits(binary_cls_pred.squeeze(1),binary_target)loss[3]0.2*binary_loss四、yolo-tensorrt接收训练的.pt使用export 接口转换为onnx使用tensorrt 的 bin工具中的trttexec将 onnx 转换为 trt 格式使用tensorrt 的 c API接口创建 m_context 和m_engine对yolo进行正常的推理以上三种操作都是常规的操作这里不进行详细的展开不会操作的可以问下AI这里只说下怎么在训练阶段增加分支然后在部署阶段 拿到增加的分支的。tensorrt中输入和输出都是在Binding中的因此通过一下代码输入Binding的个数以及每个bingding的名称、维度for(inti0;ithis-m_engine-getNbBindings();i){std::couti : this-m_engine-getBindingName(i)std::endl;printDims(this-m_engine-getBindingDimensions(i));}我转换后的trt的结果为 i0name:images 维度4x3x640x640 i1name:726维度4x576 i2name:output0 维度4x8x8400 我的类别是4 所以448images 和 output0 是在没有增加第三个分支的时候已经存在了和它们一样创建内存来接收 name: 726的数据output0的内存地址float*m_output_src_device;m_output_src_devicenullptr;checkRuntime(cudaMalloc(m_output_src_device,m_param.batch_size*m_output_area*sizeof(float)));checkRuntime(cudaFree(m_output_src_device));output0的内存地址float*m_output_src726_device;intm_output_src726_channel;m_output_src726_devicenullptr;m_output_src726_channel576;checkRuntime(cudaMalloc(m_output_src_device,m_param.batch_size*m_output_src726_channel*sizeof(float)));checkRuntime(cudaFree(m_output_src726_device));float*bindings[]{m_input_hwc_device,m_output_src_device}boolcontextm_context-executeV2((void**)bindings)修改为float*bindings[]{m_input_hwc_device,m_output_src726_device,m_output_src_device}boolcontextm_context-executeV2((void**)bindings)
