Multiple questions about DGCNN semantic segmentation

[QuanticDisaster]

I would like to implement semantic segmentation on points clouds and am trying various things, however I don't grasp the details of the implementation of the DGCNN used in the example.

For the moment, I am using GCN implementation which seems to work and is especially fast compared to DGCNN

class GCN(torch.nn.Module):
    def __init__(self, hidden_channels):
        super(GCN, self).__init__()
        torch.manual_seed(12345)
        self.conv1 = GCNConv(NUM_FEATURES, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, NUM_CLASSES)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = x.relu()
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index)
        return x

model = GCN(hidden_channels=16)#.to(device)
print(model)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = torch.nn.CrossEntropyLoss(weight=weights)

However it seems to me it fails to grasp the geometric features of neighboors and just use the features provided (for instance for solar panels, without features and only pos it seems to fail to classify while a simple threshold on the elevation which is part of the pos, should be enough to differentiate solar panels and ground), which seems normal as it is more oriented towards graph representations where position does not play an important role like social media for instance. Hence I can't use it with just the pos and no features.

However I was thinking DGCNN would be able to, and following the implementation in examples

def MLP(channels, batch_norm=True):
    return Seq(*[
        Seq(Lin(channels[i - 1], channels[i]), ReLU(), BN(channels[i]))
        for i in range(1, len(channels))
    ])

class Net(torch.nn.Module):
    def __init__(self, out_channels, k=20, aggr='max'):
        super(Net, self).__init__()

        self.conv1 = DynamicEdgeConv(MLP([2 * (NUM_FEATURES + 3), 64, 64]), k, aggr, num_workers =-1)    #WARNING : DynamicEdgeConv(MLP([2 * X... , X est égal à 3 (XYZ) + la somme des features
        self.conv2 = DynamicEdgeConv(MLP([2 * 64, 64, 64]), k, aggr, num_workers =-1)
        self.conv3 = DynamicEdgeConv(MLP([2 * 64, 64, 64]), k, aggr, num_workers =-1)
        self.lin1 = MLP([3 * 64, 1024])

        self.mlp = Seq(MLP([1024, 256]), Dropout(0.5), MLP([256, 128]),
                       Dropout(0.5), Lin(128, out_channels))

    def forward(self, data):
        x, pos, batch = data.x, data.pos, None#data.batch
        x0 = torch.cat([x, pos], dim=-1)
        x1 = self.conv1(x0, batch)
        x2 = self.conv2(x1, batch)
        x3 = self.conv3(x2, batch)
        out = self.lin1(torch.cat([x1, x2, x3], dim=1))
        out = self.mlp(out)
        return F.log_softmax(out, dim=1)

device = torch.device('cpu')
model = Net(NUM_CLASSES, k =30).to(device) 
print(model)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.8)

with no features however, I get an error saying :

Traceback (most recent call last):

  File "<ipython-input-5-99e0e19b6eb8>", line 150, in <module>
    train()

  File "<ipython-input-5-99e0e19b6eb8>", line 55, in train
    out = model(data)

  File "C:\Users\Tidop\anaconda3\envs\torch_env_38\lib\site-packages\torch\nn\modules\module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)

  File "<ipython-input-5-99e0e19b6eb8>", line 27, in forward
    x0 = torch.cat([x, pos], dim=-1)

TypeError: expected Tensor as element 0 in argument 0, but got NoneType

Using a "false" undiscriminative matrix of ones for each nodes does work but I am not sure of the result
Moreover DGCNN seems quite slow in comparison to GCN, on this type of batch :
Data(edge_index=[2, 114393], pos=[3957, 3], x=[3957, 3], y=[3957])

Each batch takes 1.5-2 second to complete and I have 406624 points alias 150 batches to compute... Is there a way to speed it up ?

There is a bit of everything as you can see I am a bit confused. Are GCN and DGCNN good choices at all for what I want to do ? DGCNN should be able to classify ground and solar panels based on their altitude (ie with only the positions and no features) contrary to GCN, right ?

Plus some more questions :
I don't understand what is the data.batch element of the implementation, what does it do, is it an int ? I am using my own data (not my own dataset, I just used ClusterData and ClusterLoad).
As we work with neighbours, the value should depends on the density of point otherwise we might fail to grasp geometric patterns if the BBox of the neighboorhood is smaller than the caracteristic geometric patterns do we agree ?

Thx in advance

[rusty1s]

Yes, your are right. GCN cannot take the spatial relation between nodes into account. A common workaround is to simply concatenate node features and positions into a single node feature matrix x, but this comes with downsides as your model now depends on absolute node positions.

The DGCNN is a good alternative, but as you said, is much slower since it will create a new graph for every layer (that is the expensive part). Another slow-down is the BatchNorm call, which will be fixed in PyTorch 1.9.0.

To fix your error, you can simple pass pos to the first layer in case x is None. A faster alternative is the PointConv layer, which simply utilizes the pre-defined edge connectivity matrix edge_index for message propagation.

We use data.batch to distinguish nodes in different graphs. In your case, you can simply drop it as your working on a single graph.

[QuanticDisaster]

Thx for your reply, it clarifies things for me.

Following your advice, I tried the pointnet2 implementation given in the examples, however while modificating the code to drop the data.batch part, I get stuck on the global_max_pool operation of the pointnet2_classification code (necessary for the semantic segmentation) :

Here are the modified parts in the forward functions and the original source code : https://github.com/rusty1s/pytorch_geometric/blob/master/examples/pointnet2_classification.py

class SAModule(torch.nn.Module):
    def __init__(self, ratio, r, nn):
        super(SAModule, self).__init__()
        self.ratio = ratio
        self.r = r
        self.conv = PointConv(nn, add_self_loops=False)
    
    def forward(self, x, pos):
        idx = fps(pos, ratio=self.ratio)
        row, col = radius(pos, pos[idx], self.r,
                          max_num_neighbors=64)
        edge_index = torch.stack([col, row], dim=0)
        x = self.conv(x, (pos, pos[idx]), edge_index)
        pos = pos[idx]
        return x, pos


class GlobalSAModule(torch.nn.Module):
    def __init__(self, nn):
        super(GlobalSAModule, self).__init__()
        self.nn = nn
    
    def forward(self, x, pos):
        
        x = self.nn(torch.cat([x, pos], dim=1))
        x = x.max()
        pos = pos.new_zeros((x.size(0), 3))   #GIVES ERROR
        return x, pos

The doc https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.glob.global_max_pool and source code show it calls torch_scatter.scatter but I am unable to understand how I should replace global_max_pool

[rusty1s]

As far as I see, you do not want to perform global_max_pool as you want to learn representations for each node/point. Therefore, I suggest to simply replace the DynamicEdgeConv calls with PointConv calls, e.g.:

self.conv1 = PointConv(MLP([NUM_FEATURES + 3, 64, 64]))

x1 = self.conv1(x, pos, edge_index)

[QuanticDisaster]

Thank you for your proposition, indeed the execution is faster.

The results were really poor at first, a 3 layer conv seemed to get stuck (or have difficulties to converge to another possibility) putting all the test point in the same class. A 2 PointConv layers network worked better, which surprised me a bit

Why use only PointConv instead of the PointNet architecture ? I assume the latter was designed specifically to be performing well on semantic segmentation so why use only a part of it ?

[rusty1s]

Yes, you can also try the complete PointNet++ segmentation model, but I just assumed that it's easier to start with a simple baseline :)

Not sure why a 2-layer net performs better though :(

[QuanticDisaster]

Well it was not wrong to assume i was better off starting easy I think :)

thx for your patience and explanations !