def forward(self, data):
batch_obj = Batch()
x, pos, batch = data.x, data.pos, data.batch
if self._precompute_multi_scale:
idx = getattr(data, "idx_{}".format(self._index), None)
else:
idx = self.sampler(pos, batch)
batch_obj.idx = idx
ms_x = []
for scale_idx in range(self.neighbour_finder.num_scales):
if self._precompute_multi_scale:
edge_index = getattr(
data, "edge_index_{}_{}".format(self._index, scale_idx),
None)
else:
row, col = self.neighbour_finder(
pos,
pos[idx],
batch_x=batch,
batch_y=batch[idx],
scale_idx=scale_idx,
)
edge_index = torch.stack([col, row], dim=0)
ms_x.append(self.conv(x, (pos, pos[idx]), edge_index, batch))
batch_obj.x = torch.cat(ms_x, -1)
batch_obj.pos = pos[idx]
batch_obj.batch = batch[idx]
copy_from_to(data, batch_obj)
return batch_obj
def forward(self, data, **kwargs):
batch_obj = Batch()
x, pos, batch = data.x, data.pos, data.batch
idx = self.sampler(pos, batch)
row, col = self.neighbour_finder(pos, pos[idx], batch_x=batch, batch_y=batch[idx])
edge_index = torch.stack([col, row], dim=0)
batch_obj.idx = idx
batch_obj.edge_index = edge_index
batch_obj.x = self.conv(x, (pos[idx], pos), edge_index, batch)
batch_obj.pos = pos[idx]
batch_obj.batch = batch[idx]
copy_from_to(data, batch_obj)
return batch_obj
def forward(self, data, **kwargs):
batch_obj = Batch()
x, pos, batch = data.x, data.pos, data.batch
idx = self.sampler(pos, batch)
batch_obj.idx = idx
ms_x = []
for scale_idx in range(self.neighbour_finder.num_scales):
row, col = self.neighbour_finder(pos, pos[idx], batch_x=batch, batch_y=batch[idx], scale_idx=scale_idx,)
edge_index = torch.stack([col, row], dim=0)
ms_x.append(self.conv(x, (pos, pos[idx]), edge_index, batch))
batch_obj.x = torch.cat(ms_x, -1)
batch_obj.pos = pos[idx]
batch_obj.batch = batch[idx]
copy_from_to(data, batch_obj)
return batch_obj
def forward(self, data):
batch_obj = Batch()
x, pos, batch = data.x, data.pos, data.batch
if self._precompute_multi_scale:
idx = getattr(data, "index_{}".format(self._index), None)
edge_index = getattr(data, "edge_index_{}".format(self._index),
None)
else:
idx = self.sampler(pos, batch)
row, col = self.neighbour_finder(pos,
pos[idx],
batch_x=batch,
batch_y=batch[idx])
edge_index = torch.stack([col, row], dim=0)
batch_obj.idx = idx
batch_obj.edge_index = edge_index
batch_obj.x = self.conv(x, (pos, pos[idx]), edge_index, batch)
batch_obj.pos = pos[idx]
batch_obj.batch = batch[idx]
copy_from_to(data, batch_obj)
return batch_obj
