Spaces:
Sleeping
Sleeping
File size: 4,949 Bytes
5b557cf | 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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | import torch
from torch_scatter import scatter
def avgPoolKernel(x,edge_index,selections,cluster,kernel_size=2,even_dirs=[0,1,7,8]):
is_even = kernel_size % 2 == 0
full_passes = kernel_size//2 - int(is_even)
# Assumes the lowest number node index is the topleft most position in the cluster
indices = torch.arange(len(x)).to(x.device)
# Find the minimum node index in each cluster and select those x values
picks = scatter(indices, cluster, dim=0, reduce='min')
# Send max pool messages the appropriate number of times
for _ in range(full_passes):
message = x[edge_index[1]]
x = scatter(message,edge_index[0],dim=0,reduce='mean') # Aggregate
# Even kernel_sizes are not symetric and are lopsided towards the bottom right corner
# Repeat the process one more time going to the bottom right
if is_even:
# Prefilter edge_index
keep = torch.zeros_like(selections,dtype=torch.bool).to(x.device)
for i in even_dirs:
keep[torch.where(selections == i)] = True
even_edge_index = edge_index[:,torch.where(keep)[0]]
message = x[even_edge_index[1]]
x = scatter(message,even_edge_index[0],dim=0,reduce='mean') # Aggregate
# Take the previously selected nodes
x = x[picks]
return x
def maxPoolKernel(x,edge_index,selections,cluster,kernel_size=2,even_dirs=[0,1,7,8]):
is_even = kernel_size % 2 == 0
full_passes = kernel_size//2 - int(is_even)
# Assumes the lowest number node index is the topleft most position in the cluster
indices = torch.arange(len(x)).to(x.device)
# Find the minimum node index in each cluster and select those x values
picks = scatter(indices, cluster, dim=0, reduce='min')
# Send max pool messages the appropriate number of times
for _ in range(full_passes):
message = x[edge_index[1]]
x = scatter(message,edge_index[0],dim=0,reduce='max') # Aggregate
# Even kernel_sizes are not symetric and are lopsided towards the bottom right corner
# Repeat the process one more time going to the bottom right
if is_even:
# Prefilter edge_index
keep = torch.zeros_like(selections,dtype=torch.bool).to(x.device)
for i in even_dirs:
keep[torch.where(selections == i)] = True
even_edge_index = edge_index[:,torch.where(keep)[0]]
message = x[even_edge_index[1]]
x = scatter(message,even_edge_index[0],dim=0,reduce='max') # Aggregate
# Take the previously selected nodes
x = x[picks]
return x
def stridePoolCluster(x,cluster):
# Assumes the lowest number node index is the topleft most position in the cluster
indices = torch.arange(len(x)).to(x.device)
# Find the minimum node index in each cluster and select those x values
picks = scatter(indices, cluster, dim=0, reduce='min')
x = x[picks]
return x
def maxPoolCluster(x,cluster):
x = scatter(x, cluster, dim=0, reduce='max')
return x
def avgPoolCluster(x,cluster,edge_index=None, edge_weight=None):
x = scatter(x, cluster, dim=0, reduce='mean')
return x
def unpoolInterpolated(x,cluster,up_edge_index,up_interps=None):
if up_interps is None:
return unpoolEdgeAverage(x,cluster,up_edge_index)
# Determine node averages based on based on interps
target_clusters = cluster[up_edge_index[1]]
node_vals = x[target_clusters]*up_interps.unsqueeze(1)
x = scatter(node_vals,up_edge_index[0],dim=0,reduce='add')
norm = scatter(up_interps,up_edge_index[0],dim=0)
x/=norm.unsqueeze(1)
return x
def unpoolBilinear(x,cluster,up_edge_index,up_selections,selection_dirs=[0,1,7,8]):
# Remove edges that won't be used for the bilinear interpolation calculation
keep = torch.zeros_like(up_selections,dtype=torch.bool).to(x.device)
for i in selection_dirs:
keep[torch.where(up_selections == i)] = True
ref_edge_index = up_edge_index[:,torch.where(keep)[0]]
cluster_index = torch.vstack((ref_edge_index[0],cluster[ref_edge_index[1]]))
cluster_index = torch.unique(cluster_index,dim=1)
x = scatter(x[cluster_index[1]],cluster_index[0],dim=0,reduce='mean')
return x
def unpoolEdgeAverage(x,cluster,up_edge_index,weighted=True):
# Interpolates based on the number of connections to each previous cluster. Works best with dense data.
# If weighted = False, clusters are weighted equally regardless of the number of connections
if weighted:
target_clusters = cluster[up_edge_index[1]]
x = scatter(x[target_clusters],up_edge_index[0],dim=0,reduce='mean')
else:
cluster_index = torch.vstack((up_edge_index[0],cluster[up_edge_index[1]]))
cluster_index = torch.unique(cluster_index,dim=1)
x = scatter(x[cluster_index[1]],cluster_index[0],dim=0,reduce='mean')
return x
def unpoolCluster(x,cluster):
return x[cluster]
|