| 63 | |
| 64 | class PVConv(nn.Module): |
| 65 | def __init__(self, in_channels, out_channels, kernel_size, resolution, attention=False, |
| 66 | dropout=0.1, with_se=False, with_se_relu=False, normalize=True, eps=0): |
| 67 | super().__init__() |
| 68 | self.in_channels = in_channels |
| 69 | self.out_channels = out_channels |
| 70 | self.kernel_size = kernel_size |
| 71 | self.resolution = resolution |
| 72 | |
| 73 | self.voxelization = Voxelization(resolution, normalize=normalize, eps=eps) |
| 74 | voxel_layers = [ |
| 75 | nn.Conv3d(in_channels, out_channels, kernel_size, stride=1, padding=kernel_size // 2), |
| 76 | nn.GroupNorm(num_groups=8, num_channels=out_channels), |
| 77 | Swish() |
| 78 | ] |
| 79 | voxel_layers += [nn.Dropout(dropout)] if dropout is not None else [] |
| 80 | voxel_layers += [ |
| 81 | nn.Conv3d(out_channels, out_channels, kernel_size, stride=1, padding=kernel_size // 2), |
| 82 | nn.GroupNorm(num_groups=8, num_channels=out_channels), |
| 83 | Attention(out_channels, 8) if attention else Swish() |
| 84 | ] |
| 85 | if with_se: |
| 86 | voxel_layers.append(SE3d(out_channels, use_relu=with_se_relu)) |
| 87 | self.voxel_layers = nn.Sequential(*voxel_layers) |
| 88 | self.point_features = SharedMLP(in_channels, out_channels) |
| 89 | |
| 90 | def forward(self, inputs): |
| 91 | features, coords, temb = inputs |