下载
中文
注册

beta)torch_npu.contrib.function.npu_bbox_coder_encode_yolo

接口原型

torch_npu.contrib.function.npu_bbox_coder_encode_yolo(bboxes, gt_bboxes, stride):

功能描述

使用NPU OP获取将bbox转换为gt_bbox的框回归转换deltas。

参数说明

  • bboxes (torch.Tensor) - 源框,例如锚点框。支持数据类型:float、half。
  • gt_bboxes (torch.Tensor) - 转换目标框,例如真值框。支持数据类型:float、half。
  • stride (torch.Tensor) - Bbox步长。仅支持int张量。

输出说明

torch.Tensor - 框转换deltas。

支持的型号

  • Atlas 训练系列产品
  • Atlas A2 训练系列产品
  • Atlas 推理系列产品

调用示例

>>> from torch_npu.contrib.function import npu_bbox_coder_encode_yolo
>>> A = 1024
>>> bboxes = torch.randint(0, 512, size=(A, 4)).npu()
>>> gt_bboxes = torch.randint(0, 512, size=(A, 4)).npu()
>>> stride = torch.randint(0, 32, size=(A,)).npu()
>>> out = npu_bbox_coder_encode_yolo(bboxes, gt_bboxes, stride)
>>> torch.npu.synchronize()
>>> print('_npu_bbox_coder_encode_yolo done. output shape is ', out.shape)