| 16 | super(SparseGroupNorm, self).__init__(num_groups, num_channels, eps, affine) |
| 17 | |
| 18 | def forward(self, input: SparseTensor) -> SparseTensor: |
| 19 | nfeats = torch.zeros_like(input.feats) |
| 20 | for k in range(input.shape[0]): |
| 21 | if DEBUG: |
| 22 | assert (input.coords[input.layout[k], 0] == k).all(), f"SparseGroupNorm: batch index mismatch" |
| 23 | bfeats = input.feats[input.layout[k]] |
| 24 | bfeats = bfeats.permute(1, 0).reshape(1, input.shape[1], -1) |
| 25 | bfeats = super().forward(bfeats) |
| 26 | bfeats = bfeats.reshape(input.shape[1], -1).permute(1, 0) |
| 27 | nfeats[input.layout[k]] = bfeats |
| 28 | return input.replace(nfeats) |
| 29 | |
| 30 | |
| 31 | class SparseLayerNorm(nn.LayerNorm): |