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