DAF算子(DeformableAggregationFunction)的转onnx

DAF算子是用cuda实现的。转onnx的时候,如果数据和模型都在cpu上,那么在DAF算子这一步会报错说发现张量不是在同一个设备上,有的cpu,有的gpu。因此需要在进入这个算子的时候,先把数据和模型搬到gpu上,在DAF算子算完之后,把它返回的结果搬回cpu上。用到的hook DAF算子如下:

@staticmethod
def Hook_DeformableAggregationFunction_forward(
    ctx,
    mc_ms_feat,
    spatial_shape,
    scale_start_index,
    sampling_location,
    weights,
    roundoff,
):
    mc_ms_feat = mc_ms_feat.contiguous().float()
    spatial_shape = spatial_shape.contiguous().int()
    scale_start_index = scale_start_index.contiguous().int()
    sampling_location = sampling_location.contiguous().float()
    weights = weights.contiguous().float()
    roundoff = torch.tensor(roundoff, requires_grad=False).to(weights.device).contiguous().float()
    output = deformable_aggregation_ext.deformable_aggregation_forward(
        mc_ms_feat.cuda(),
        spatial_shape.cuda(),
        scale_start_index.cuda(),
        sampling_location.cuda(),
        weights.cuda(),
        roundoff.cuda(),
    )
    output = output / roundoff
    return output.to(mc_ms_feat.device)

  

posted @ 2026-06-09 17:07  Picassooo  阅读(11)  评论(0)    收藏  举报