Adapted models for image segmentation

[kimzed]

I am having troubles finding a good architecture / data preprocessing for an image segmentation task for my research.
I used the slic algorithm to compute super pixels, and i have a dataset with fully connected graphs and another with edges between touching nodes.

with this settings i tried different graph layers (GCN, a more basic GNN, an attention layer) but it is difficult to get good results compared to a very simple CNN.

I am going towards adding more nearest neighbour because apparently a fully connected graph can oversmoothed information during the aggregation phase, and too few edges is not optimal for message passing.

In terms of architecture though it is difficult to get a good one, does someone have experience with it or know a good architecture for this task ? (or a simple article with a github repo)

Here is one of my models as an example:

from torch_geometric.nn import GraphConv
import torch
import torch.nn.functional as F
import torch.nn as nn
import utils as functions
from torch.nn import Linear

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class Gnn(torch.nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(Gnn, self).__init__()
        torch.manual_seed(42)
        self.learning_rate = 0.01
        self.initial_layer = GraphConv(input_dim, hidden_dim)
        self.layer1 = GraphConv(hidden_dim, hidden_dim)
        self.layer2 = GraphConv(hidden_dim, hidden_dim)
        self.out = Linear(hidden_dim, output_dim)

    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch

        x = self.initial_layer(x=x, edge_index=edge_index)
        x = torch.tanh(x)
        x = F.dropout(x, training=self.training)
        x = self.layer1(x=x, edge_index=edge_index)
        x = torch.tanh(x)
        x = self.layer2(x=x, edge_index=edge_index)
        x = torch.tanh(x)
        x = self.out(x)

        return x

[rusty1s]

Importantly, you have to add node positions to your graph, either as node features or relative edge features between nodes. As such, please try the following

x = torch.cat([data.x, data.pos], dim=-1)

inside your model.

[rusty1s]

What's the output of data.pos.min(dim=0) and data.pos.max(dim=0). You may need to consider normalizing the positions.

[kimzed]

I get the following, i do not understand what indices are though

torch.return_types.min(
values=tensor([1.1087, 1.1304], device='cuda:0'),
indices=tensor([323, 644], device='cuda:0'))

torch.return_types.max(
values=tensor([36.1071, 36.0000], device='cuda:0'),
indices=tensor([ 59, 727], device='cuda:0'))

[kimzed]

Do you recommend a z-score normalization for node features? when using CNN I just subtract the mean channel-wise, it is usually more adapted for my data. but maybe it is not adapted for a graph model. I also compute distance between nodes and it is not normalized

[kimzed]

after normalizing the coordinates the model is able to converge again. so indeed normalization helped for the coordinates

[rusty1s]

Super :)