-
Notifications
You must be signed in to change notification settings - Fork 1
/
pool.py
36 lines (32 loc) · 1.27 KB
/
pool.py
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
from torch import Tensor
from torch_scatter import segment_csr
from torch_geometric.utils import degree
import torch
def global_add_pool(x: Tensor, batch: Tensor):
sorted, indices = batch.sort(stable=True)
x = x[indices.long()]
deg = degree(sorted)
ptr = deg.cumsum(dim=0)
ptr = torch.cat([torch.tensor([0], device=x.device), ptr], dim=0).long()
return segment_csr(x, ptr, reduce="sum")
def global_mean_pool(x: Tensor, batch: Tensor):
sorted, indices = batch.sort(stable=True)
x = x[indices.long()]
deg = degree(sorted)
ptr = deg.cumsum(dim=0)
ptr = torch.cat([torch.tensor([0], device=x.device), ptr], dim=0).long()
return segment_csr(x, ptr, reduce="mean")
def global_max_pool(x: Tensor, batch: Tensor):
sorted, indices = batch.sort(stable=True)
x = x[indices.long()]
deg = degree(sorted)
ptr = deg.cumsum(dim=0)
ptr = torch.cat([torch.tensor([0], device=x.device), ptr], dim=0).long()
return segment_csr(x, ptr, reduce="max")
def global_min_pool(x: Tensor, batch: Tensor):
sorted, indices = batch.sort(stable=True)
x = x[indices.long()]
deg = degree(sorted)
ptr = deg.cumsum(dim=0)
ptr = torch.cat([torch.tensor([0], device=x.device), ptr], dim=0).long()
return segment_csr(x, ptr, reduce="min")