def __call__(self, data, norm=True):
row, col = data.edge_index
N = data.num_nodes
deg = degree(row, N, dtype=torch.float)
if norm:
deg = deg / deg.max()
deg_col = deg[col]
min_deg, _ = scatter_min(deg_col, row, dim_size=N)
min_deg[min_deg > 10000] = 0
max_deg, _ = scatter_max(deg_col, row, dim_size=N)
max_deg[max_deg < -10000] = 0
mean_deg = scatter_mean(deg_col, row, dim_size=N)
std_deg = scatter_std(deg_col, row, dim_size=N)
x = torch.stack([deg, min_deg, max_deg, mean_deg, std_deg], dim=1)
if data.x is not None:
data.x = data.x.view(-1, 1) if data.x.dim() == 1 else data.x
data.x = torch.cat([data.x, x], dim=-1)
else:
data.x = x
return data
def test_std(dtype, device, bias):
src = tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]], dtype, device)
index = tensor([[0, 0, 0, 0, 0], [1, 1, 1, 1, 1]], torch.long, device)
out = scatter_std(src, index, dim=-1, unbiased=bias)
expected = src.std(dim=-1, unbiased=bias)
assert out.tolist() == [[expected[0].item(), 0], [0, expected[0].item()]]
def test_empty_std(dtype, device):
out = torch.zeros(1, 5, dtype=dtype, device=device)
src = tensor([], dtype, device).view(0, 5)
index = tensor([], torch.long, device).view(0, 5)
out = scatter_std(src, index, dim=0, out=out)
assert out.tolist() == [[0, 0, 0, 0, 0]]
def aggregate(self,
inputs: Tensor,
index: Tensor,
dim_size: Optional[int] = None) -> Tensor:
if self.aggr == 'softmax':
out = scatter_softmax(inputs * self.t, index, dim=self.node_dim)
return scatter(inputs * out,
index,
dim=self.node_dim,
dim_size=dim_size,
reduce='sum')
elif self.aggr == 'softmax_sg':
out = scatter_softmax(inputs * self.t, index,
dim=self.node_dim).detach()
return scatter(inputs * out,
index,
dim=self.node_dim,
dim_size=dim_size,
reduce='sum')
elif self.aggr == 'stat':
_mean = scatter_mean(inputs,
index,
dim=self.node_dim,
dim_size=dim_size)
_std = scatter_std(inputs,
index,
dim=self.node_dim,
dim_size=dim_size).detach()
_min = scatter_min(inputs,
index,
dim=self.node_dim,
dim_size=dim_size)[0]
_max = scatter_max(inputs,
index,
dim=self.node_dim,
dim_size=dim_size)[0]
_mean = _mean.unsqueeze(dim=-1)
_std = _std.unsqueeze(dim=-1)
_min = _min.unsqueeze(dim=-1)
_max = _max.unsqueeze(dim=-1)
stat = torch.cat([_mean, _std, _min, _max], dim=-1)
stat = self.lin_stat(stat)
stat = stat.squeeze(dim=-1)
return stat
else:
min_value, max_value = 1e-7, 1e1
torch.clamp_(inputs, min_value, max_value)
out = scatter(torch.pow(inputs, self.p),
index,
dim=self.node_dim,
dim_size=dim_size,
reduce='mean')
torch.clamp_(out, min_value, max_value)
return torch.pow(out, 1 / self.p)
def test_std(dtype, device, bias):
src = tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]], dtype, device)
index = tensor([[0, 0, 0, 0, 0], [1, 1, 1, 1, 1]], torch.long, device)
out = scatter_std(src, index, dim=-1, unbiased=bias)
std = src.std(dim=-1, unbiased=bias)[0].item()
expected = tensor([[std, 0], [0, std]], dtype, device)
assert torch.allclose(out, expected)
def test_std():
src = torch.tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]], dtype=torch.float)
index = torch.tensor([[0, 0, 0, 0, 0], [1, 1, 1, 1, 1]], dtype=torch.long)
out = scatter_std(src, index, dim=-1, unbiased=True)
std = src.std(dim=-1, unbiased=True)[0]
expected = torch.tensor([[std, 0], [0, std]])
assert torch.allclose(out, expected)
def test_std():
src = torch.tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]], dtype=torch.float)
src.requires_grad_()
index = torch.tensor([[0, 0, 0, 0, 0], [1, 1, 1, 1, 1]], dtype=torch.long)
out = scatter_std(src, index, dim=-1, unbiased=True)
std = src.std(dim=-1, unbiased=True)[0]
expected = torch.tensor([[std, 0], [0, std]])
assert torch.allclose(out, expected)
out.backward(torch.randn_like(out))
def __call__(self, data):
row, col = data.edge_index
N = data.num_nodes
deg = degree(row, N, dtype=torch.float)
deg_col = deg[col]
value = 1e16
min_deg, _ = scatter_min(deg_col, row, dim_size=N, fill_value=value)
min_deg[min_deg == value] = 0
max_deg, _ = scatter_max(deg_col, row, dim_size=N)
mean_deg = scatter_mean(deg_col, row, dim_size=N)
std_deg = scatter_std(deg_col, row, dim_size=N)
x = torch.stack([deg, min_deg, max_deg, mean_deg, std_deg], dim=1)
if data.x is not None:
data.x = data.x.view(-1, 1) if data.x.dim() == 1 else data.x
data.x = torch.cat([data.x, x], dim=-1)
else:
data.x = x
return data
src = torch.Tensor([[2, 1, 1, 4, 2], [1, 2, 1, 2, 4]]).float() index = torch.tensor([[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]]) out = src.new_ones((2, 6)) out = scatter_div(src, index, out=out) print(out) # tensor([[1.0000, 1.0000, 0.2500, 0.5000, 0.5000, 1.0000], # [0.5000, 0.2500, 0.5000, 1.0000, 1.0000, 1.0000]]) # 最大最小平均值 src = torch.Tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]]) index = torch.tensor([[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]]) out, argmax = scatter_max(src, index) print(out, argmax) out, argmin = scatter_min(src, index) print(out, argmin) out = scatter_mean(src, index) print(out) out = scatter_mul(src, index) print(out) out = scatter_std(src, index) print(out) out = scatter_sub(src, index) print(out)
