| 142 | |
| 143 | |
| 144 | class VNBatchNorm(nn.Module): |
| 145 | def __init__(self, num_features, dim): |
| 146 | super(VNBatchNorm, self).__init__() |
| 147 | self.dim = dim |
| 148 | if dim == 3 or dim == 4: |
| 149 | self.bn = nn.BatchNorm1d(num_features) |
| 150 | elif dim == 5: |
| 151 | self.bn = nn.BatchNorm2d(num_features) |
| 152 | |
| 153 | def forward(self, x): |
| 154 | ''' |
| 155 | x: point features of shape [B, N_feat, 3, N_samples, ...] |
| 156 | ''' |
| 157 | # norm = torch.sqrt((x*x).sum(2)) |
| 158 | norm = torch.norm(x, dim=2) + EPS |
| 159 | norm_bn = self.bn(norm) |
| 160 | norm = norm.unsqueeze(2) |
| 161 | norm_bn = norm_bn.unsqueeze(2) |
| 162 | x = x / norm * norm_bn |
| 163 | |
| 164 | return x |
| 165 | |
| 166 | |
| 167 | class VNMaxPool(nn.Module): |