1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
|
import re
from typing import List, Optional, Union
import torch
from torch import Tensor
import torch_geometric
from torch_geometric.data import Data
from torch_geometric.data.datapipes import functional_transform
from torch_geometric.transforms import BaseTransform
from torch_geometric.utils import one_hot, scatter
@functional_transform('grid_sampling')
class GridSampling(BaseTransform):
r"""Clusters points into fixed-sized voxels
(functional name: :obj:`grid_sampling`).
Each cluster returned is a new point based on the mean of all points
inside the given cluster.
Args:
size (float or [float] or Tensor): Size of a voxel (in each dimension).
start (float or [float] or Tensor, optional): Start coordinates of the
grid (in each dimension). If set to :obj:`None`, will be set to the
minimum coordinates found in :obj:`data.pos`.
(default: :obj:`None`)
end (float or [float] or Tensor, optional): End coordinates of the grid
(in each dimension). If set to :obj:`None`, will be set to the
maximum coordinates found in :obj:`data.pos`.
(default: :obj:`None`)
"""
def __init__(
self,
size: Union[float, List[float], Tensor],
start: Optional[Union[float, List[float], Tensor]] = None,
end: Optional[Union[float, List[float], Tensor]] = None,
) -> None:
self.size = size
self.start = start
self.end = end
def forward(self, data: Data) -> Data:
num_nodes = data.num_nodes
assert data.pos is not None
c = torch_geometric.nn.voxel_grid(data.pos, self.size, data.batch,
self.start, self.end)
c, perm = torch_geometric.nn.pool.consecutive.consecutive_cluster(c)
for key, item in data.items():
if bool(re.search('edge', key)):
raise ValueError(f"'{self.__class__.__name__}' does not "
f"support coarsening of edges")
if torch.is_tensor(item) and item.size(0) == num_nodes:
if key == 'y':
item = scatter(one_hot(item), c, dim=0, reduce='sum')
data[key] = item.argmax(dim=-1)
elif key == 'batch':
data[key] = item[perm]
else:
data[key] = scatter(item, c, dim=0, reduce='mean')
return data
def __repr__(self) -> str:
return f'{self.__class__.__name__}(size={self.size})'
|
