def forward(self, data):
data.x = F.elu(
self.bn1(self.conv0(data.x, data.edge_index, data.edge_attr)))
cluster = voxel_grid(data.pos, data.batch, size=[4, 3])
data = max_pool(cluster, data, transform=T.Cartesian(cat=False))
data = self.conv1(data)
cluster = voxel_grid(data.pos, data.batch, size=[16, 12])
data = max_pool(cluster, data, transform=T.Cartesian(cat=False))
data = self.conv2(data)
cluster = voxel_grid(data.pos, data.batch, size=[30, 23])
data = max_pool(cluster, data, transform=T.Cartesian(cat=False))
data = self.conv3(data)
cluster = voxel_grid(data.pos, data.batch, size=[60, 45])
x = max_pool_x(cluster, data.x, data.batch, size=16)
# x = max_pool_x(cluster, data.x, data.batch)
x = x[0].view(-1, self.fc1.weight.size(1))
x = self.fc1(x)
x = F.elu(x)
x = self.bn(x)
x = self.drop_out(x)
x = self.fc2(x)
return F.log_softmax(x, dim=1)
def forward(self, data):
data.x = F.elu(self.conv1(data.x, data.edge_index))
data.x = self.bn1(data.x)
cluster = voxel_grid(data.pos, data.batch, size=[4,4])
data = max_pool(cluster, data, transform=T.Cartesian(cat=False))
data.x = F.elu(self.conv2(data.x, data.edge_index))
data.x = self.bn2(data.x)
cluster = voxel_grid(data.pos, data.batch, size=[6,6])
data = max_pool(cluster, data, transform=T.Cartesian(cat=False))
data.x = F.elu(self.conv3(data.x, data.edge_index))
data.x = self.bn3(data.x)
cluster = voxel_grid(data.pos, data.batch, size=[20,20])
data = max_pool(cluster, data, transform=T.Cartesian(cat=False))
data.x = F.elu(self.conv4(data.x, data.edge_index))
data.x = self.bn4(data.x)
cluster = voxel_grid(data.pos, data.batch, size=[32,32])
x = max_pool_x(cluster, data.x, data.batch, size=32)
x = x.view(-1, self.fc1.weight.size(1))
x = F.elu(self.fc1(x))
x = F.dropout(x, training=self.training)
x = self.fc2(x)
return F.log_softmax(x, dim=1)
def forward(self, data):
data.x = F.elu(
self.bn1(self.conv1(data.x, data.edge_index, data.edge_attr)))
cluster = voxel_grid(data.pos, data.batch, size=4)
data = max_pool(cluster, data, transform=T.Cartesian(cat=False))
data = self.block1(data)
cluster = voxel_grid(data.pos, data.batch, size=6)
data = max_pool(cluster, data, transform=T.Cartesian(cat=False))
data = self.block2(data)
cluster = voxel_grid(data.pos, data.batch, size=24)
data = max_pool(cluster, data, transform=T.Cartesian(cat=False))
data = self.block3(data)
cluster = voxel_grid(data.pos, data.batch, size=64)
x = max_pool_x(cluster, data.x, data.batch, size=8)
# if your torch-geometric version is below 1.3.2(roughly, we do not test all versions), use x.view() instead of x[0].view()
# x = x.view(-1, self.fc1.weight.size(1))
x = x[0].view(-1, self.fc1.weight.size(1))
x = self.fc1(x)
x = F.elu(x)
x = self.bn(x)
x = self.drop_out(x)
x = self.fc2(x)
return F.log_softmax(x, dim=1)
def test_voxel_grid():
pos = torch.Tensor([[0, 0], [11, 9], [2, 8], [2, 2], [8, 3]])
batch = torch.tensor([0, 0, 0, 1, 1])
assert voxel_grid(pos, batch, size=5).tolist() == [0, 5, 3, 6, 7]
cluster = voxel_grid(pos, batch, size=5, start=-1, end=[18, 14])
assert cluster.tolist() == [0, 10, 4, 16, 17]
def forward(self, data):
# 下采样，得到新的数据
cluster = nn_geometric.voxel_grid(
data.pos,
data.batch,
self.pool_rad, # self.pool_rad=0.1，找到目标池化后数据的范围
start=data.pos.min(dim=0)[0] - self.pool_rad * 0.5,
end=data.pos.max(dim=0)[0] + self.pool_rad * 0.5)
cluster, perm = consecutive_cluster(cluster)
new_pos = scatter_('mean', data.pos, cluster)
new_batch = data.batch[perm]
cluster = nearest(data.pos, new_pos, data.batch, new_batch)
data.x = scatter_(
self._aggr, data.x, cluster,
dim_size=new_pos.size(0)) # 根据self._aggr判断是什么池化，最大还是平均
data.pos = new_pos
data.batch = new_batch
data.edge_attr = None
data = self.graph_reg(data) # 下采样完成后，再进行构图操作
return data
def forward(self, data):
cluster = nn_geometric.voxel_grid(
data.pos,
data.batch,
self.pool_rad,
start=data.pos.min(dim=0)[0] - self.pool_rad * 0.5,
end=data.pos.max(dim=0)[0] + self.pool_rad * 0.5)
cluster, perm = consecutive_cluster(cluster)
new_pos = scatter(data.pos, cluster, dim=0, reduce='mean')
new_batch = data.batch[perm]
cluster = nearest(data.pos, new_pos, data.batch, new_batch)
cluster, perm = consecutive_cluster(cluster)
data.x = scatter(data.x, cluster, dim=0, reduce=self._aggr)
data.pos = scatter(data.pos, cluster, dim=0, reduce='mean')
data.batch = data.batch[perm]
data.edge_index = None
data.edge_attr = None
data = self.graph_reg(data)
return data
def __call__(self, data):
num_nodes = data.num_nodes
if "batch" not in data:
batch = data.pos.new_zeros(num_nodes, dtype=torch.long)
else:
batch = data.batch
cluster = voxel_grid(data.pos, batch, self.size, self.start, self.end)
cluster, perm = consecutive_cluster(cluster)
for key, item in data:
if bool(re.search("edge", key)):
raise ValueError(
"GridSampling does not support coarsening of edges")
if torch.is_tensor(item) and item.size(0) == num_nodes:
if key == "y":
item = F.one_hot(item, num_classes=self.num_classes)
item = scatter_add(item, cluster, dim=0)
data[key] = item.argmax(dim=-1)
elif key == "batch":
data[key] = item[perm]
else:
data[key] = scatter_mean(item, cluster, dim=0)
return data
def _process(self, data):
if self._mode == "last":
data = shuffle_data(data)
coords = ((data.pos) / self._grid_size).int()
if "batch" not in data:
cluster = grid_cluster(coords, torch.tensor([1, 1, 1]))
else:
cluster = voxel_grid(coords, data.batch, 1)
cluster, unique_pos_indices = consecutive_cluster(cluster)
skip_keys = []
if self._quantize_coords:
skip_keys.append("pos")
data = group_data(data,
cluster,
unique_pos_indices,
mode=self._mode,
skip_keys=skip_keys)
if self._quantize_coords:
data.pos = coords[unique_pos_indices]
return data
def test_single_voxel_grid():
pos = torch.Tensor([[0, 0], [1, 1], [2, 2], [3, 3], [4, 4]])
edge_index = torch.tensor([[0, 0, 3], [1, 2, 4]])
batch = torch.tensor([0, 0, 0, 1, 1])
x = torch.randn(5, 16)
cluster = voxel_grid(pos, size=5, batch=batch)
assert cluster.tolist() == [0, 0, 0, 1, 1]
data = Batch(x=x, edge_index=edge_index, pos=pos, batch=batch)
data = avg_pool(cluster, data)
cluster_no_batch = voxel_grid(pos, size=5)
assert cluster_no_batch.tolist() == [0, 0, 0, 0, 0]
data_no_batch = Batch(x=x, edge_index=edge_index, pos=pos)
data_no_batch = avg_pool(cluster_no_batch, data_no_batch)
def torch_subsampling(data, subsampling_parameter):
batch = torch.zeros(data.shape[0], dtype=torch.long)
pool = voxel_grid(torch.tensor(data[:, :3]), batch, subsampling_parameter)
points = pool_pos(pool, torch.tensor(data[:, :3]))
normals = pool_pos(pool, torch.tensor(data[:, 3:]))
del pool
del batch
return points.numpy(), normals.numpy()
def forward(self, data):
data.x = F.elu(self.conv1(data.x, data.edge_index, data.edge_attr))
cluster = voxel_grid(data.pos, data.batch, size=5, start=0, end=28)
data = max_pool(cluster, data, transform=transform)
data.x = F.elu(self.conv2(data.x, data.edge_index, data.edge_attr))
cluster = voxel_grid(data.pos, data.batch, size=7, start=0, end=28)
data = max_pool(cluster, data, transform=transform)
data.x = F.elu(self.conv3(data.x, data.edge_index, data.edge_attr))
cluster = voxel_grid(data.pos, data.batch, size=14, start=0, end=27.99)
x = max_pool_x(cluster, data.x, data.batch, size=4)
x = x.view(-1, self.fc1.weight.size(1))
x = F.elu(self.fc1(x))
x = F.dropout(x, training=self.training)
x = self.fc2(x)
return F.log_softmax(x, dim=1)
def _prepare_data(self, data):
coords = torch.round((data.pos) / self._grid_size).long()
cluster = voxel_grid(coords, data.batch, 1)
cluster, unique_pos_indices = consecutive_cluster(cluster)
coords = coords[unique_pos_indices]
new_batch = data.batch[unique_pos_indices]
new_pos = data.pos[unique_pos_indices]
x = self._aggregate(data.x, cluster, unique_pos_indices)
sparse_data = Batch(x=x, pos=new_pos, coords=coords, batch=new_batch)
return sparse_data, cluster
def sample(self, pos=None, x=None, batch=None):
if len(pos.shape) != 2:
raise ValueError("This class is for sparse data and expects the pos tensor to be of dimension 2")
pool = voxel_grid(pos, batch, self._subsampling_param)
pool, perm = consecutive_cluster(pool)
batch = pool_batch(perm, batch)
if x is not None:
return pool_pos(pool, x), pool_pos(pool, pos), batch
else:
return None, pool_pos(pool, pos), batch
def forward(self, data):
data = self.conv1(data)
cluster = voxel_grid(data.pos, data.batch, size=2)
data = max_pool(cluster, data, transform=T.Cartesian(cat=False))
data = self.conv2(data)
cluster = voxel_grid(data.pos, data.batch, size=4)
data = max_pool(cluster, data, transform=T.Cartesian(cat=False))
data = self.conv3(data)
cluster = voxel_grid(data.pos, data.batch, size=7)
x = max_pool_x(cluster, data.x, data.batch, size=25)
# x = max_pool_x(cluster, data.x, data.batch)
x = x[0].view(-1, self.fc1.weight.size(1))
x = self.fc1(x)
x = F.elu(x)
x = self.bn(x)
x = self.drop_out(x)
x = self.fc2(x)
return F.log_softmax(x, dim=1)
def _process(self, data):
num_nodes = data.num_nodes
if "batch" not in data:
batch = data.pos.new_zeros(num_nodes, dtype=torch.long)
else:
batch = data.batch
cluster = voxel_grid(data.pos, batch, self.size, self.start, self.end)
cluster, perm = consecutive_cluster(cluster)
return group_data(data, cluster, perm, mode="mean")
def _process(self, data):
if self._mode == "last":
data = shuffle_data(data)
coords = torch.round((data.pos) / self._grid_size)
if "batch" not in data:
cluster = grid_cluster(coords, torch.tensor([1, 1, 1]))
else:
cluster = voxel_grid(coords, data.batch, 1)
cluster, unique_pos_indices = consecutive_cluster(cluster)
data = group_data(data, cluster, unique_pos_indices, mode=self._mode)
if self._quantize_coords:
data.coords = coords[unique_pos_indices].int()
return data
def forward(self, x):
data = x
data.x = F.relu(
self.bn1(self.conv1(data.x, data.edge_index, data.edge_attr)))
data.x = self.layer1(data)
data.x = self.layer2(data.x, data.edge_index, data.edge_attr)
data.x = self.layer3(data.x, data.edge_index, data.edge_attr)
data.x = self.layer4(data.x, data.edge_index, data.edge_attr)
# clustering in Spline COnv
cluster = voxel_grid(data.pos, data.batch, size=4)
x = max_pool_x(cluster, data.x, data.batch, size=4)
x = x.view(-1, self.fc.weight.size(1))
x = self.fc(x)
return x
def __call__(self, data):
num_nodes = data.num_nodes
if 'batch' not in data:
batch = data.pos.new_zeros(num_nodes, dtype=torch.long)
else:
batch = data.batch
cluster = voxel_grid(data.pos, batch, self.size, self.start, self.end)
cluster, perm = consecutive_cluster(cluster)
for key, item in data:
if bool(re.search('edge', key)):
raise ValueError(
'GridSampling does not support coarsening of edges')
if torch.is_tensor(item) and item.size(0) == num_nodes:
if key == 'batch':
data[key] = item[perm]
else:
data[key] = scatter_mean(item, cluster, dim=0)
return data
def __call__(self, data):
try:
batch = data.batch
except AttributeError:
batch = torch.zeros(data.x.shape[0])
# points = [data.pos]
# First downsample
points = [data.pos]
list_pool = []
list_batch = [batch]
for ind, voxel_size in enumerate(self.list_voxel_size):
pool = voxel_grid(points[-1], list_batch[-1], voxel_size)
pool, perm = consecutive_cluster(pool)
list_batch.append(pool_batch(perm, list_batch[-1]))
points.append(pool_pos(pool, points[-1]))
try:
res = MultiScaleBatch(
batch=data.batch,
list_pool=list_pool,
points=points,
list_batch=list_batch,
x=data.x,
y=data.y,
pos=data.pos)
except AttributeError:
res = MultiScaleData(
list_pool=list_pool,
points=points,
x=data.x,
y=data.y,
pos=data.pos)
return res
def __call__(self, data):
num_nodes = data.num_nodes
pos = data.pos
batch = data.batch
pool = voxel_grid(pos, batch, self._subsampling_param)
pool, _ = consecutive_cluster(pool)
for key, item in data:
if bool(re.search('edge', key)):
continue
if torch.is_tensor(item) and item.size(0) == num_nodes:
if key == 'y':
one_hot = torch.zeros((item.shape[0], self._num_classes))\
.scatter(1, item.unsqueeze(-1), 1)
aggr_labels = scatter_add(one_hot, pool, dim=0)
data[key] = torch.argmax(aggr_labels, -1)
else:
data[key] = pool_pos(pool, item).to(item.dtype)
return data
