Пример #1
0
    def roi_align(g, input, rois, spatial_scale, pooled_height, pooled_width,
                  sampling_ratio, aligned):
        batch_indices = _cast_Long(
            g,
            squeeze(
                g,
                select(
                    g, rois, 1,
                    g.op('Constant',
                         value_t=torch.tensor([0], dtype=torch.long))), 1),
            False)
        rois = select(
            g, rois, 1,
            g.op('Constant',
                 value_t=torch.tensor([1, 2, 3, 4], dtype=torch.long)))
        if aligned:
            warnings.warn(
                "ONNX export of ROIAlign with aligned=True does not match PyTorch when using malformed boxes,"
                " ONNX forces ROIs to be 1x1 or larger.")
            scale = torch.tensor(0.5 / spatial_scale).to(dtype=torch.float)
            rois = g.op("Sub", rois, scale)

        # ONNX doesn't support negative sampling_ratio
        if sampling_ratio < 0:
            warnings.warn("ONNX doesn't support negative sampling ratio,"
                          "therefore is is set to 0 in order to be exported.")
            sampling_ratio = 0
        return g.op('RoiAlign',
                    input,
                    rois,
                    batch_indices,
                    spatial_scale_f=spatial_scale,
                    output_height_i=pooled_height,
                    output_width_i=pooled_width,
                    sampling_ratio_i=sampling_ratio)
Пример #2
0
 def symbolic_multi_label_nms(g, boxes, scores, iou_threshold):
     boxes = unsqueeze(g, boxes, 0)
     scores = unsqueeze(g, unsqueeze(g, scores, 0), 0)
     max_output_per_class = g.op("Constant",
                                 value_t=torch.tensor([sys.maxsize],
                                                      dtype=torch.long))
     iou_threshold = g.op("Constant",
                          value_t=torch.tensor([iou_threshold],
                                               dtype=torch.float))
     nms_out = g.op("NonMaxSuppression", boxes, scores,
                    max_output_per_class, iou_threshold)
     return squeeze(
         g,
         select(
             g, nms_out, 1,
             g.op("Constant", value_t=torch.tensor([2], dtype=torch.long))),
         1)
Пример #3
0
    def roi_align(g, input, rois, spatial_scale, pooled_height, pooled_width,
                  sampling_ratio, aligned):
        batch_indices = _cast_Long(
            g,
            squeeze(
                g,
                select(
                    g, rois, 1,
                    g.op("Constant",
                         value_t=torch.tensor([0], dtype=torch.long))), 1),
            False)
        rois = select(
            g, rois, 1,
            g.op("Constant",
                 value_t=torch.tensor([1, 2, 3, 4], dtype=torch.long)))
        # TODO: Remove this warning after ONNX opset 16 is supported.
        if aligned:
            warnings.warn(
                "ROIAlign with aligned=True is not supported in ONNX, but will be supported in opset 16. "
                "The workaround is that the user need apply the patch "
                "https://github.com/microsoft/onnxruntime/pull/8564 "
                "and build ONNXRuntime from source.")

        # ONNX doesn't support negative sampling_ratio
        if sampling_ratio < 0:
            warnings.warn(
                "ONNX doesn't support negative sampling ratio, therefore is set to 0 in order to be exported."
            )
            sampling_ratio = 0
        return g.op(
            "RoiAlign",
            input,
            rois,
            batch_indices,
            spatial_scale_f=spatial_scale,
            output_height_i=pooled_height,
            output_width_i=pooled_width,
            sampling_ratio_i=sampling_ratio,
        )