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)

浙公网安备 33010602011771号