| import torch.nn as nn |
|
|
| from . import functional as F |
| from .voxelization import Voxelization |
| from .shared_mlp import SharedMLP |
| import torch |
|
|
| __all__ = ['PVConv'] |
|
|
|
|
| class PVConv(nn.Module): |
| def __init__( |
| self, in_channels, out_channels, kernel_size, resolution, with_se=False, normalize=True, eps=0, scale_pvcnn=False, |
| device='cuda'): |
| super().__init__() |
| self.in_channels = in_channels |
| self.out_channels = out_channels |
| self.kernel_size = kernel_size |
| self.resolution = resolution |
| self.voxelization = Voxelization(resolution, normalize=normalize, eps=eps, scale_pvcnn=scale_pvcnn) |
| voxel_layers = [ |
| nn.Conv3d(in_channels, out_channels, kernel_size, stride=1, padding=kernel_size // 2, device=device), |
| nn.InstanceNorm3d(out_channels, eps=1e-4, device=device), |
| nn.LeakyReLU(0.1, True), |
| nn.Conv3d(out_channels, out_channels, kernel_size, stride=1, padding=kernel_size // 2, device=device), |
| nn.InstanceNorm3d(out_channels, eps=1e-4, device=device), |
| nn.LeakyReLU(0.1, True), |
| ] |
| self.voxel_layers = nn.Sequential(*voxel_layers) |
| self.point_features = SharedMLP(in_channels, out_channels, device=device) |
|
|
| def forward(self, inputs): |
| features, coords = inputs |
| voxel_features, voxel_coords = self.voxelization(features, coords) |
| voxel_features = self.voxel_layers(voxel_features) |
| devoxel_features = F.trilinear_devoxelize(voxel_features, voxel_coords, self.resolution, self.training) |
| fused_features = devoxel_features + self.point_features(features) |
| return fused_features, coords, voxel_features |
|
|
