(self, input: SparseTensor)
| 20 | self.factor = tuple(factor) if isinstance(factor, (list, tuple)) else factor |
| 21 | |
| 22 | def forward(self, input: SparseTensor) -> SparseTensor: |
| 23 | DIM = input.coords.shape[-1] - 1 |
| 24 | factor = self.factor if isinstance(self.factor, tuple) else (self.factor,) * DIM |
| 25 | assert DIM == len(factor), 'Input coordinates must have the same dimension as the downsample factor.' |
| 26 | |
| 27 | coord = list(input.coords.unbind(dim=-1)) |
| 28 | for i, f in enumerate(factor): |
| 29 | coord[i+1] = coord[i+1] // f |
| 30 | |
| 31 | MAX = [coord[i+1].max().item() + 1 for i in range(DIM)] |
| 32 | OFFSET = torch.cumprod(torch.tensor(MAX[::-1]), 0).tolist()[::-1] + [1] |
| 33 | code = sum([c * o for c, o in zip(coord, OFFSET)]) |
| 34 | code, idx = code.unique(return_inverse=True) |
| 35 | |
| 36 | new_feats = torch.scatter_reduce( |
| 37 | torch.zeros(code.shape[0], input.feats.shape[1], device=input.feats.device, dtype=input.feats.dtype), |
| 38 | dim=0, |
| 39 | index=idx.unsqueeze(1).expand(-1, input.feats.shape[1]), |
| 40 | src=input.feats, |
| 41 | reduce='mean' |
| 42 | ) |
| 43 | new_coords = torch.stack( |
| 44 | [code // OFFSET[0]] + |
| 45 | [(code // OFFSET[i+1]) % MAX[i] for i in range(DIM)], |
| 46 | dim=-1 |
| 47 | ) |
| 48 | out = SparseTensor(new_feats, new_coords, input.shape,) |
| 49 | out._scale = tuple([s // f for s, f in zip(input._scale, factor)]) |
| 50 | out._spatial_cache = input._spatial_cache |
| 51 | |
| 52 | out.register_spatial_cache(f'upsample_{factor}_coords', input.coords) |
| 53 | out.register_spatial_cache(f'upsample_{factor}_layout', input.layout) |
| 54 | out.register_spatial_cache(f'upsample_{factor}_idx', idx) |
| 55 | |
| 56 | return out |
| 57 | |
| 58 | |
| 59 | class SparseUpsample(nn.Module): |
nothing calls this directly
no test coverage detected