def forward(self, data):
data.x = self.datanorm * data.x
data.x = self.inputnet(data.x)
data.edge_index = to_undirected(
knn_graph(data.x,
self.k,
data.batch,
loop=False,
flow=self.edgeconv1.flow))
data.x = self.edgeconv1(data.x, data.edge_index)
weight = normalized_cut_2d(data.edge_index, data.x)
cluster = graclus(data.edge_index, weight, data.x.size(0))
data.edge_attr = None
data = max_pool(cluster, data)
data.edge_index = to_undirected(
knn_graph(data.x,
self.k,
data.batch,
loop=False,
flow=self.edgeconv2.flow))
data.x = self.edgeconv2(data.x, data.edge_index)
weight = normalized_cut_2d(data.edge_index, data.x)
cluster = graclus(data.edge_index, weight, data.x.size(0))
x, batch = max_pool_x(cluster, data.x, data.batch)
x = global_max_pool(x, batch)
return self.output(x).squeeze(-1)
def forward(self, data):
# Use the coords for the first knn step
print('data.x:', data.x.size())
print('data.batch:', data.batch.size())
clustering1 = to_undirected(
knn_graph(data.x,
self.k,
data.batch,
loop=False,
flow=self.edgeconv1.flow))
print('clustering1:', clustering1.size())
out1 = self.edgeconv1(data.features, clustering1)
print('out1:', out1.size())
raise Exception('stop')
# Now use the outputted features of the previous layer for the knn
clustering2 = to_undirected(
knn_graph(out1,
self.k,
data.batch,
loop=False,
flow=self.edgeconv2.flow))
out2 = self.edgeconv2(out1, clustering2)
clustering3 = to_undirected(
knn_graph(out2,
self.k,
data.batch,
loop=False,
flow=self.edgeconv3.flow))
out3 = self.edgeconv3(out2, clustering3)
# Cat all outputs together
edgeconv_out = torch.cat([data.features, out1, out2, out3])
# Run the output layer
return self.output(edgeconv_out).squeeze(-1)
def __call__(self, data):
def is_in(ar1, ar2):
return (ar1[..., None] == ar2).any(-1)
print('add mesh structure for shape', data.shape_id.numpy())
face_t = data.face.t()
edge_index_t = data.edge_index.t()
face_edges_t = torch.tensor([[x for x in range(len(edge_index_t)) if
is_in(edge_index_t[x], f).all() and edge_index_t[x][0] <
edge_index_t[x][1]] for f in face_t],
device=torch.device('cuda'))
data.face_edges = face_edges_t.t()
gt_edge_index = torch.cat([data.gt_face[:2], data.gt_face[1:], data.gt_face[::2]], dim=1)
gt_edge_index = to_undirected(gt_edge_index, num_nodes=data.num_nodes)
data.gt_edge_index = gt_edge_index
return data
def process(self):
with open(self.raw_paths[0], 'r') as f:
data = f.read().split('\n')[1:-1]
x = [[float(v) for v in r.split('\t')[1].split(',')] for r in data]
x = torch.tensor(x, dtype=torch.float)
y = [int(r.split('\t')[2]) for r in data]
y = torch.tensor(y, dtype=torch.long)
with open(self.raw_paths[1], 'r') as f:
data = f.read().split('\n')[1:-1]
data = [[int(v) for v in r.split('\t')] for r in data]
edge_index = torch.tensor(data, dtype=torch.long).t().contiguous()
edge_index = to_undirected(edge_index)
edge_index, _ = coalesce(edge_index, None, x.size(0), x.size(0))
data = Data(x=x, edge_index=edge_index, y=y)
data = data if self.pre_transform is None else self.pre_transform(data)
torch.save(self.collate([data]), self.processed_paths[0])