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