Spaces:
Running
on
A10G
Running
on
A10G
File size: 4,256 Bytes
251e479 |
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 |
# Copyright (c) OpenMMLab. All rights reserved.
import torch
from torch import nn as nn
from torch.autograd import Function
import annotator.uniformer.mmcv as mmcv
from ..utils import ext_loader
ext_module = ext_loader.load_ext(
'_ext', ['roiaware_pool3d_forward', 'roiaware_pool3d_backward'])
class RoIAwarePool3d(nn.Module):
"""Encode the geometry-specific features of each 3D proposal.
Please refer to `PartA2 <https://arxiv.org/pdf/1907.03670.pdf>`_ for more
details.
Args:
out_size (int or tuple): The size of output features. n or
[n1, n2, n3].
max_pts_per_voxel (int, optional): The maximum number of points per
voxel. Default: 128.
mode (str, optional): Pooling method of RoIAware, 'max' or 'avg'.
Default: 'max'.
"""
def __init__(self, out_size, max_pts_per_voxel=128, mode='max'):
super().__init__()
self.out_size = out_size
self.max_pts_per_voxel = max_pts_per_voxel
assert mode in ['max', 'avg']
pool_mapping = {'max': 0, 'avg': 1}
self.mode = pool_mapping[mode]
def forward(self, rois, pts, pts_feature):
"""
Args:
rois (torch.Tensor): [N, 7], in LiDAR coordinate,
(x, y, z) is the bottom center of rois.
pts (torch.Tensor): [npoints, 3], coordinates of input points.
pts_feature (torch.Tensor): [npoints, C], features of input points.
Returns:
pooled_features (torch.Tensor): [N, out_x, out_y, out_z, C]
"""
return RoIAwarePool3dFunction.apply(rois, pts, pts_feature,
self.out_size,
self.max_pts_per_voxel, self.mode)
class RoIAwarePool3dFunction(Function):
@staticmethod
def forward(ctx, rois, pts, pts_feature, out_size, max_pts_per_voxel,
mode):
"""
Args:
rois (torch.Tensor): [N, 7], in LiDAR coordinate,
(x, y, z) is the bottom center of rois.
pts (torch.Tensor): [npoints, 3], coordinates of input points.
pts_feature (torch.Tensor): [npoints, C], features of input points.
out_size (int or tuple): The size of output features. n or
[n1, n2, n3].
max_pts_per_voxel (int): The maximum number of points per voxel.
Default: 128.
mode (int): Pooling method of RoIAware, 0 (max pool) or 1 (average
pool).
Returns:
pooled_features (torch.Tensor): [N, out_x, out_y, out_z, C], output
pooled features.
"""
if isinstance(out_size, int):
out_x = out_y = out_z = out_size
else:
assert len(out_size) == 3
assert mmcv.is_tuple_of(out_size, int)
out_x, out_y, out_z = out_size
num_rois = rois.shape[0]
num_channels = pts_feature.shape[-1]
num_pts = pts.shape[0]
pooled_features = pts_feature.new_zeros(
(num_rois, out_x, out_y, out_z, num_channels))
argmax = pts_feature.new_zeros(
(num_rois, out_x, out_y, out_z, num_channels), dtype=torch.int)
pts_idx_of_voxels = pts_feature.new_zeros(
(num_rois, out_x, out_y, out_z, max_pts_per_voxel),
dtype=torch.int)
ext_module.roiaware_pool3d_forward(rois, pts, pts_feature, argmax,
pts_idx_of_voxels, pooled_features,
mode)
ctx.roiaware_pool3d_for_backward = (pts_idx_of_voxels, argmax, mode,
num_pts, num_channels)
return pooled_features
@staticmethod
def backward(ctx, grad_out):
ret = ctx.roiaware_pool3d_for_backward
pts_idx_of_voxels, argmax, mode, num_pts, num_channels = ret
grad_in = grad_out.new_zeros((num_pts, num_channels))
ext_module.roiaware_pool3d_backward(pts_idx_of_voxels, argmax,
grad_out.contiguous(), grad_in,
mode)
return None, None, grad_in, None, None, None
|