Spaces:
Sleeping
Sleeping
| """Sparse inference primitives extracted for the released SPAC VQ-VAE.""" | |
| from typing import * | |
| import torch | |
| import torch.nn as nn | |
| from . import SparseTensor | |
| class SparseDownsample(nn.Module): | |
| def __init__(self, factor: Union[int, Tuple[int, ...], List[int]], mode='mean'): | |
| super(SparseDownsample, self).__init__() | |
| self.factor = tuple(factor) if isinstance(factor, (list, tuple)) else factor | |
| self.mode = mode | |
| def forward(self, input: SparseTensor) -> SparseTensor: | |
| DIM = input.coords.shape[-1] - 1 | |
| factor = self.factor if isinstance(self.factor, tuple) else (self.factor,) * DIM | |
| assert DIM == len(factor), 'Input coordinates must have the same dimension as the downsample factor.' | |
| coord = list(input.coords.unbind(dim=-1)) | |
| for (i, f) in enumerate(factor): | |
| coord[i + 1] = coord[i + 1] // f | |
| MAX = [coord[i + 1].max().item() + 1 for i in range(DIM)] | |
| OFFSET = torch.cumprod(torch.tensor(MAX[::-1]), 0).tolist()[::-1] + [1] | |
| code = sum([c * o for (c, o) in zip(coord, OFFSET)]) | |
| (code, idx) = code.unique(return_inverse=True) | |
| dtype = input.feats.dtype | |
| new_feats = torch.scatter_reduce(torch.zeros(code.shape[0], input.feats.shape[1], device=input.feats.device, dtype=torch.float64), dim=0, index=idx.unsqueeze(1).expand(-1, input.feats.shape[1]), src=input.feats.double(), reduce=self.mode) | |
| new_feats = new_feats.to(dtype) | |
| new_coords = torch.stack([code // OFFSET[0]] + [code // OFFSET[i + 1] % MAX[i] for i in range(DIM)], dim=-1) | |
| out = SparseTensor(new_feats, new_coords, input.shape) | |
| out._scale = tuple([s // f for (s, f) in zip(input._scale, factor)]) | |
| out._spatial_cache = input._spatial_cache | |
| out.register_spatial_cache(f'upsample_{factor}_coords', input.coords) | |
| out.register_spatial_cache(f'upsample_{factor}_layout', input.layout) | |
| out.register_spatial_cache(f'upsample_{factor}_idx', idx) | |
| return out | |
| class SparseSubdivide(nn.Module): | |
| def __init__(self): | |
| super(SparseSubdivide, self).__init__() | |
| def forward(self, input: SparseTensor) -> SparseTensor: | |
| DIM = input.coords.shape[-1] - 1 | |
| n_cube = torch.ones([2] * DIM, device=input.device, dtype=torch.int) | |
| n_coords = torch.nonzero(n_cube) | |
| n_coords = torch.cat([torch.zeros_like(n_coords[:, :1]), n_coords], dim=-1) | |
| factor = n_coords.shape[0] | |
| assert factor == 2 ** DIM | |
| new_coords = input.coords.clone() | |
| new_coords[:, 1:] *= 2 | |
| new_coords = new_coords.unsqueeze(1) + n_coords.unsqueeze(0).to(new_coords.dtype) | |
| new_feats = input.feats.unsqueeze(1).expand(input.feats.shape[0], factor, *input.feats.shape[1:]) | |
| out = SparseTensor(new_feats.flatten(0, 1), new_coords.flatten(0, 1), input.shape) | |
| out._scale = input._scale * 2 | |
| out._spatial_cache = input._spatial_cache | |
| return out | |