def forward(self, data):
data.x = F.elu(self.conv1(data.x, data.edge_index, data.edge_attr))
weight = normalized_cut_2d(data.edge_index, data.pos)
cluster = graclus(data.edge_index, weight, data.x.size(0))
data.edge_attr = None
data = max_pool(cluster, data, transform=transform)
data.x = F.elu(self.conv2(data.x, data.edge_index, data.edge_attr))
weight = normalized_cut_2d(data.edge_index, data.pos)
cluster = graclus(data.edge_index, weight, data.x.size(0))
x, batch = max_pool_x(cluster, data.x, data.batch)
#x = global_mean_pool(x, batch)
x_min = torch_scatter.scatter_min(x, batch, dim=0)[0]
gather_idxs = batch.expand(x.shape[1], -1).t()
gather_mins = torch.gather(x_min, 0, gather_idxs)
s = F.relu(-gather_mins)
x = x + s
x = self.aggregator(x, batch)
s_out = self.aggregator(s, batch)
x = x - s_out
x = F.elu(self.fc1(x))
x = F.dropout(x, training=self.training)
return F.log_softmax(self.fc2(x), dim=1)
def forward(self, data):
x, edge_index, batch = data.x, data.edge_index, data.batch
if self.encode_edge:
x = self.atom_encoder(x)
x = self.conv1(x, edge_index, data.edge_attr)
else:
x = self.conv1(x, edge_index)
x = F.relu(x)
xs = [global_mean_pool(x, batch)]
for i, conv in enumerate(self.convs):
x = F.relu(conv(x, edge_index))
xs += [global_mean_pool(x, batch)]
if self.pooling_type != 'none':
if self.pooling_type == 'complement':
complement = batched_negative_edges(edge_index=edge_index, batch=batch, force_undirected=True)
cluster = graclus(complement, num_nodes=x.size(0))
elif self.pooling_type == 'graclus':
cluster = graclus(edge_index, num_nodes=x.size(0))
data = Batch(x=x, edge_index=edge_index, batch=batch)
data = max_pool(cluster, data)
x, edge_index, batch = data.x, data.edge_index, data.batch
if not self.no_cat:
x = self.jump(xs)
else:
x = global_mean_pool(x, batch)
x = F.relu(self.lin1(x))
x = self.lin2(x)
return x
def forward(self, data):
data.x = F.elu(self.conv1a(data.x, data.edge_index, data.edge_attr))
data.x = F.elu(self.conv1b(data.x, data.edge_index, data.edge_attr))
# data.x = F.elu(self.conv1c(data.x, data.edge_index, data.edge_attr))
# data.x = self.bn1(data.x)
weight = normalized_cut_2d(data.edge_index, data.pos)
cluster1 = graclus(data.edge_index, weight, data.x.size(0))
pos1 = data.pos
edge_index1 = data.edge_index
batch1 = data.batch if hasattr(data, 'batch') else None
# weights1 = bweights(data, cluster1)
data = max_pool(cluster1, data, transform=T.Cartesian(cat=False))
data.x = F.elu(self.conv2a(data.x, data.edge_index, data.edge_attr))
data.x = F.elu(self.conv2b(data.x, data.edge_index, data.edge_attr))
# data.x = F.elu(self.conv2c(data.x, data.edge_index, data.edge_attr))
weight = normalized_cut_2d(data.edge_index, data.pos)
cluster2 = graclus(data.edge_index, weight, data.x.size(0))
pos2 = data.pos
edge_index2 = data.edge_index
batch2 = data.batch if hasattr(data, 'batch') else None
# weights2 = bweights(data, cluster2)
data = max_pool(cluster2, data, transform=T.Cartesian(cat=False))
# upsample
# data = recover_grid_barycentric(data, weights=weights2, pos=pos2, edge_index=edge_index2, cluster=cluster2,
# batch=batch2, transform=T.Cartesian(cat=False))
data.x = F.elu(self.conv3a(data.x, data.edge_index, data.edge_attr))
data.x = F.elu(self.conv3b(data.x, data.edge_index, data.edge_attr))
data = recover_grid(data,
pos2,
edge_index2,
cluster2,
batch=batch2,
transform=T.Cartesian(cat=False))
# data = recover_grid_barycentric(data, weights=weights1, pos=pos1, edge_index=edge_index1, cluster=cluster1,
# batch=batch1, transform=T.Cartesian(cat=False))
data.x = F.elu(self.conv4a(data.x, data.edge_index, data.edge_attr))
data.x = F.elu(self.conv4b(data.x, data.edge_index, data.edge_attr))
data = recover_grid(data,
pos1,
edge_index1,
cluster1,
batch=batch1,
transform=T.Cartesian(cat=False))
# TODO handle contract on trainer and evaluator
data.x = F.elu(self.convout(data.x, data.edge_index, data.edge_attr))
x = data.x
# return F.sigmoid(x)
return x
def forward(self, data):
data.x = F.elu(self.conv1(data.x, data.edge_index, data.edge_attr))
cluster = graclus(
data.edge_index,
torch.reshape(data.edge_attr, (data.edge_attr.shape[0], )),
data.x.size(0))
data = max_pool(cluster, data)
data.x = F.elu(self.conv2(data.x, data.edge_index, data.edge_attr))
cluster = graclus(
data.edge_index,
torch.reshape(data.edge_attr, (data.edge_attr.shape[0], )),
data.x.size(0))
x, batch = max_pool_x(cluster, data.x, data.batch)
x = global_mean_pool(x, batch)
x = F.elu(self.fc1(x))
x = F.dropout(x, training=self.training)
return F.log_softmax(self.fc2(x), dim=1)
def forward(self, data):
data.x = F.relu(self.conv1(data.x, data.edge_index))
data.x = F.dropout(data.x, training=self.training)
weight = normalized_cut_2d(data.edge_index, data.pos)
cluster = graclus(data.edge_index, weight, data.x.size(0))
data.x, batch = max_pool_x(cluster, data.x, data.batch)
data.x = global_mean_pool(data.x, batch)
data.x = self.fc1(data.x)
data.x = F.log_softmax(data.x, dim=1)
return data.x
def forward(self, data):
data.x = F.elu(self.conv1(data.x, data.edge_index, data.edge_attr))
######## calculate similarity between nodes
weight = normalized_cut_2d(data.edge_index, data.pos)
######### graph clustering without the need of eigenvector
cluster = graclus(data.edge_index, weight, data.x.size(0))
data.edge_attr = None
########## Pools and coarsens a graph. All nodes within the same cluster will be represented as one node and appply transform
data = max_pool(cluster, data, transform=transform)
########## 2nd conv net
data.x = F.elu(self.conv2(data.x, data.edge_index, data.edge_attr))
weight = normalized_cut_2d(data.edge_index, data.pos)
cluster = graclus(data.edge_index, weight, data.x.size(0))
############## Max-Pools node features according to the clustering defined in cluster
x, batch = max_pool_x(cluster, data.x, data.batch)
############## Returns batch-wise graph-level-outputs by averaging node features across the node dimension
x = global_mean_pool(x, batch)
x = F.elu(self.fc1(x))
x = F.dropout(x, training=self.training)
return F.log_softmax(self.fc2(x), dim=1)
def forward(self, data):
data.x = F.elu(self.conv1(data.x, data.edge_index, data.edge_attr))
weight = normalized_cut_2d(data.edge_index, data.pos)
cluster = graclus(data.edge_index, weight, data.x.size(0))
data.edge_attr = None
data = max_pool(cluster, data, transform=T.Cartesian(cat=False))
data.x = F.elu(self.conv2(data.x, data.edge_index, data.edge_attr))
weight = normalized_cut_2d(data.edge_index, data.pos)
cluster = graclus(data.edge_index, weight, data.x.size(0))
data = max_pool(cluster, data, transform=T.Cartesian(cat=False))
data.x = F.elu(self.conv3(data.x, data.edge_index, data.edge_attr))
x = F.elu(self.fc1(data.x))
x = F.dropout(x, training=self.training, p=self.dropout)
x = self.fc2(x)
y = global_mean_pool(x, data.batch)
if (self.wgan):
return y
return torch.sigmoid(y)
def forward(self, data):
for i in range(self.layers_num):
data.x = self.conv_layers[i](data.x, data.pos, data.edge_index)
if self.use_cluster_pooling:
weight = normalized_cut_2d(data.edge_index, data.pos)
cluster = graclus(data.edge_index, weight, data.x.size(0))
data = max_pool(cluster,
data,
transform=T.Cartesian(cat=False))
data.x = global_mean_pool(data.x, data.batch)
x = self.fc1(data.x)
return F.log_softmax(x, dim=1)
def forward(self, data):
pos, edge_index, batch = data.pos, data.edge_index, data.batch
real_batch_size = pos.size(0) / self.nr_points
real_batch_size = int(real_batch_size)
# Build first edges
edge_index = knn_graph(pos, self.k, batch, loop=False)
#extract features in 3d
_, _, features_dd, _ = self.ds1(pos, edge_index, None)
#graclus
cluster = graclus(edge_index)
pos_gra, batch_gra = avg_pool_x(cluster, pos, batch)
features_gra, _ = max_pool_x(cluster, features_dd, batch)
#knn(f)
with torch.no_grad():
edge_index_gra = knn_graph(features_gra.norm(dim=2),
self.k,
batch_gra,
loop=False)
# DD2
_, _, features_dd2, _ = self.dd2(pos_gra, edge_index_gra, features_gra)
y1 = self.nn1(features_dd2)
y1_pool, _ = max_pool_x(batch_gra, y1, batch_gra)
y1_pool = torch.nn.functional.relu(y1_pool)
y1_pool = self.bn1(y1_pool)
y2 = self.nn2(y1_pool)
y2 = torch.nn.functional.relu(y2)
y2 = self.bn2(y2)
y3 = self.nn3(y2)
y3 = torch.nn.functional.relu(y3)
y3 = self.bn3(y3)
y4 = self.nn4(y3)
out = self.sm(y4)
return out
def forward(self, data):
x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr
x = self.conv1(x)
x = F.elu(x)
weight = normalized_cut_2d(data.edge_index, data.pos)
cluster = graclus(data.edge_index, weight, x.size(0))
data.edge_attr = None
data = max_pool(cluster, data, transform=T.Cartesian(cat=False))
x = F.elu(self.conv1(x, edge_index, edge_attr))
x = self.conv2(x, edge_index, edge_attr)
x = F.elu(self.conv3(x, edge_index, edge_attr))
x = self.conv4(x, edge_index, edge_attr)
x = F.elu(self.conv5(x, edge_index, edge_attr))
x = self.conv6(x, edge_index, edge_attr)
x = F.dropout(x, training=self.training)
return F.log_softmax(x, dim=1)
def forward(self, graph):
data = graph
data.x = torch.cat([data.pos, data.x], dim=1)
for i, monet_layer in enumerate(self.monet_layers[:-1]):
data.x = F.relu(
monet_layer(data.x, data.edge_index, data.edge_attr))
weight = normalized_cut_2d(data.edge_index, data.pos)
cluster = graclus(data.edge_index, weight, data.x.size(0))
if i == 0:
data.edge_attr = None
data = max_pool(cluster, data, transform=T.Cartesian(cat=False))
data.x = self.monet_layers[-1](data.x, data.edge_index, data.edge_attr)
for linear_layer in self.linear_layers[:-1]:
x = global_mean_pool(data.x, data.batch)
x = F.relu(linear_layer(x))
x = F.dropout(x)
return F.log_softmax(self.linear_layers[-1](x), dim=1)
def test_graclus():
edge_index = torch.tensor([[0, 1], [1, 0]])
assert graclus(edge_index).tolist() == [0, 0]
def forward(self, data):
cluster = geom_nn.graclus(data.edge_index, num_nodes=data.x.size(0))
data = geom_nn.max_pool(cluster, data)
return data
def apply_graclus_pooling(x, edge_index, batch, method="max"):
cluster = graclus(edge_index)
func = max_pool if method == "max" else avg_pool
new_data = func(cluster, Batch(x=x, edge_index=edge_index, batch=batch))
return new_data.x, new_data.edge_index, new_data.batch
def mgpool(x, pos, edge_index, batch, mask=None):
adj_values = torch.ones(edge_index.shape[1]).cuda()
cluster = graclus(edge_index)
cluster, perm = consecutive_cluster(cluster)
index = torch.stack([cluster, torch.arange(0, x.shape[0]).cuda()], dim=0)
values = torch.ones(cluster.shape[0], dtype=torch.float).cuda()
uniq, inv, counts = torch.unique(cluster,
return_inverse=True,
return_counts=True)
newsize = uniq.shape[0]
origsize = x.shape[0]
new_batch = pool_batch(perm, batch)
# Compute random walk graph laplacian:
laplacian_index, laplacian_weights = get_laplacian(edge_index,
normalization='rw')
laplacian_index, laplacian_weights = torch_sparse.coalesce(
laplacian_index, laplacian_weights, m=origsize, n=origsize)
index, values = torch_sparse.coalesce(index, values, m=newsize,
n=origsize) # P^T matrix
new_feat = torch_sparse.spmm(index,
values,
m=newsize,
n=origsize,
matrix=x) # P^T X
new_pos = torch_sparse.spmm(index,
values,
m=newsize,
n=origsize,
matrix=pos) # P^T POS
new_adj, new_adj_val = torch_sparse.spspmm(index,
values,
edge_index,
adj_values,
m=newsize,
k=origsize,
n=origsize,
coalesced=True) # P^T A
index, values = torch_sparse.transpose(index,
values,
m=newsize,
n=origsize,
coalesced=True) # P
new_adj, new_adj_val = torch_sparse.spspmm(new_adj,
new_adj_val,
index,
values,
m=newsize,
k=origsize,
n=newsize,
coalesced=True) # (P^T A) P
# Precompute QP :
values = torch.ones(cluster.shape[0], dtype=torch.float).cuda()
index, values = torch_sparse.spspmm(laplacian_index,
laplacian_weights,
index,
values,
m=origsize,
k=origsize,
n=newsize,
coalesced=True)
return new_adj, new_feat, new_pos, new_batch, index, values, origsize, newsize
