| 98 | |
| 99 | |
| 100 | class PVConvReLU(nn.Module): |
| 101 | def __init__(self, in_channels, out_channels, kernel_size, resolution, attention=False, leak=0.2, |
| 102 | dropout=0.1, with_se=False, with_se_relu=False, normalize=True, eps=0): |
| 103 | super().__init__() |
| 104 | self.in_channels = in_channels |
| 105 | self.out_channels = out_channels |
| 106 | self.kernel_size = kernel_size |
| 107 | self.resolution = resolution |
| 108 | |
| 109 | self.voxelization = Voxelization(resolution, normalize=normalize, eps=eps) |
| 110 | voxel_layers = [ |
| 111 | nn.Conv3d(in_channels, out_channels, kernel_size, stride=1, padding=kernel_size // 2), |
| 112 | nn.BatchNorm3d(out_channels), |
| 113 | nn.LeakyReLU(leak, True) |
| 114 | ] |
| 115 | voxel_layers += [nn.Dropout(dropout)] if dropout is not None else [] |
| 116 | voxel_layers += [ |
| 117 | nn.Conv3d(out_channels, out_channels, kernel_size, stride=1, padding=kernel_size // 2), |
| 118 | nn.BatchNorm3d(out_channels), |
| 119 | Attention(out_channels, 8) if attention else nn.LeakyReLU(leak, True) |
| 120 | ] |
| 121 | if with_se: |
| 122 | voxel_layers.append(SE3d(out_channels, use_relu=with_se_relu)) |
| 123 | self.voxel_layers = nn.Sequential(*voxel_layers) |
| 124 | self.point_features = SharedMLP(in_channels, out_channels) |
| 125 | |
| 126 | def forward(self, inputs): |
| 127 | features, coords, temb = inputs |
| 128 | voxel_features, voxel_coords = self.voxelization(features, coords) |
| 129 | voxel_features = self.voxel_layers(voxel_features) |
| 130 | voxel_features = F.trilinear_devoxelize(voxel_features, voxel_coords, self.resolution, self.training) |
| 131 | fused_features = voxel_features + self.point_features(features) |
| 132 | return fused_features, coords, temb |
nothing calls this directly
no outgoing calls
no test coverage detected