Spaces:
Build error
Build error
File size: 4,976 Bytes
28c256d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
from typing import Tuple, Union
import torch
import torch.nn as nn
from torch.autograd import Function
from torch.autograd.function import once_differentiable
from torch.nn.modules.utils import _pair
from ..utils import ext_loader
ext_module = ext_loader.load_ext(
'_ext', ['bezier_align_forward', 'bezier_align_backward'])
class BezierAlignFunction(Function):
@staticmethod
def forward(ctx,
input: torch.Tensor,
beziers: torch.Tensor,
output_size: Union[int, Tuple[int, int]],
spatial_scale: Union[int, float] = 1.0,
sampling_ratio: int = 0,
aligned: bool = True) -> torch.Tensor:
ctx.output_size = _pair(output_size)
ctx.spatial_scale = spatial_scale
ctx.input_shape = input.size()
ctx.sampling_ratio = sampling_ratio
ctx.aligned = aligned
assert beziers.size(1) == 17
output_shape = (beziers.size(0), input.size(1), ctx.output_size[0],
ctx.output_size[1])
output = input.new_zeros(output_shape)
ext_module.bezier_align_forward(
input,
beziers,
output,
aligned_height=ctx.output_size[0],
aligned_width=ctx.output_size[1],
spatial_scale=ctx.spatial_scale,
sampling_ratio=ctx.sampling_ratio,
aligned=ctx.aligned)
ctx.save_for_backward(beziers)
return output
@staticmethod
@once_differentiable
def backward(ctx, grad_output: torch.Tensor):
beziers = ctx.saved_tensors[0]
grad_input = grad_output.new_zeros(ctx.input_shape)
grad_output = grad_output.contiguous()
ext_module.bezier_align_backward(
grad_output,
beziers,
grad_input,
aligned_height=ctx.output_size[0],
aligned_width=ctx.output_size[1],
spatial_scale=ctx.spatial_scale,
sampling_ratio=ctx.sampling_ratio,
aligned=ctx.aligned)
return grad_input, None, None, None, None, None
bezier_align = BezierAlignFunction.apply
class BezierAlign(nn.Module):
"""Bezier align pooling layer.
Args:
output_size (tuple): h, w
spatial_scale (float): scale the input boxes by this number
sampling_ratio (int): number of inputs samples to take for each
output sample. 0 to take samples densely for current models.
aligned (bool): if False, use the legacy implementation in
MMDetection. If True, align the results more perfectly.
Note:
The implementation of BezierAlign is modified from
https://github.com/aim-uofa/AdelaiDet
The meaning of aligned=True:
Given a continuous coordinate c, its two neighboring pixel
indices (in our pixel model) are computed by floor(c - 0.5) and
ceil(c - 0.5). For example, c=1.3 has pixel neighbors with discrete
indices [0] and [1] (which are sampled from the underlying signal
at continuous coordinates 0.5 and 1.5). But the original roi_align
(aligned=False) does not subtract the 0.5 when computing
neighboring pixel indices and therefore it uses pixels with a
slightly incorrect alignment (relative to our pixel model) when
performing bilinear interpolation.
With `aligned=True`,
we first appropriately scale the ROI and then shift it by -0.5
prior to calling roi_align. This produces the correct neighbors;
The difference does not make a difference to the model's
performance if ROIAlign is used together with conv layers.
"""
def __init__(
self,
output_size: Tuple,
spatial_scale: Union[int, float],
sampling_ratio: int,
aligned: bool = True,
) -> None:
super().__init__()
self.output_size = _pair(output_size)
self.spatial_scale = float(spatial_scale)
self.sampling_ratio = int(sampling_ratio)
self.aligned = aligned
def forward(self, input: torch.Tensor,
beziers: torch.Tensor) -> torch.Tensor:
"""BezierAlign forward.
Args:
inputs (Tensor): input features.
beziers (Tensor): beziers for align.
"""
return bezier_align(input, beziers, self.output_size,
self.spatial_scale, self.sampling_ratio,
self.aligned)
def __repr__(self):
s = self.__class__.__name__
s += f'(output_size={self.output_size}, '
s += f'spatial_scale={self.spatial_scale})'
s += f'sampling_ratio={self.sampling_ratio})'
s += f'aligned={self.aligned})'
return s
|