Pytorch Geometric - Train model on "PairData" (two graphs in single data object)

username: seanco92
cooked: <p>Hi guys,</p>
<p>I have dataset such each Data object is consists from two separate graphs (The graph are not connected - <strong>I don’t want them to affect each other until the linear stage</strong> of my model).</p>
<p>I built the data objects and dataset based on the “PairData” example here:<br>
<a href="https://pytorch-geometric.readthedocs.io/en/latest/notes/batching.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch-geometric.readthedocs.io/en/latest/notes/batching.html</a><br>
My Data object:</p>
<blockquote>
<p>class TwoGraphsData(Data):</p>
<pre><code>  def __init__(self, x_a=None, edge_index_a=None, edge_attr_a=None, x_b=None, edge_index_b=None,
             edge_attr_b=None, y=None, linker_size=0):
    super().__init__()
    self.x_a = x_a
    self.edge_index_a = edge_index_a
    self.edge_attr_a = edge_attr_a
    self.x_b = x_b
    self.edge_index_b = edge_index_b
    self.edge_attr_b = edge_attr_b
    self.linker_size = linker_size
</code></pre>
</blockquote>
<p>Note: The <strong>order</strong> of the two sub-graphs inside the Data object is <strong>doesn’t matter</strong>. Each sub-graph may be the ‘a’ graph or the ‘b’ graph. In fact, the model has to be order invariant.</p>
<p>My model has some GCNconv , pooling and linear layers.<br>
The <strong>forward</strong> function for <strong>single</strong> graph in <strong>regular</strong> data object is:</p>
<blockquote>
<p>x, edge_index, batch = data.x.float(), data.edge_index, data.batch<br>
edge_attr = torch.flatten(data.edge_attr)</p>
<pre><code>    x = F.relu(self.conv1(x, edge_index, edge_attr))
    x, edge_index, edge_attr, batch = self.pool1(x, edge_index, edge_attr, batch)
    x1 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)

    x = F.relu(self.conv2(x, edge_index, edge_attr))
    x, edge_index, edge_attr, batch = self.pool2(x, edge_index, edge_attr, batch)
    x2 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)

    x = F.relu(self.conv3(x, edge_index, edge_attr))
    x3 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)

    x = F.relu(x1) + F.relu(x2) + F.relu(x3)
    x = torch.cat((x, torch.reshape(data.linker, (data.linker.size()[0], 1))), dim=1)
    x = F.relu(self.lin1(x))
    x = F.dropout(x, p=self.dropout_ratio, training=self.training)
    x = F.relu(self.lin2(x))
    x = self.lin3(x)
    return x
</code></pre>
</blockquote>
<p>I want to activate the <strong>conv</strong> and <strong>pooling</strong> layers on <strong>each sub-graph separately.</strong><br>
I want to activate the <strong>linear</strong> layers on the <strong>concatenation</strong> of conv+pooling of sub-graph a and conv+pooling of sub-graph b results.</p>
<p><strong>I wonder what is the right way to handle “PairData” for correct training.</strong><br>
I thought about few options:</p>
<p><strong>a.</strong> Just use each conv and pool on x_a and x_b in a row. something like:</p>
<blockquote>
<p>x_a = F.relu(self.conv1(x_a, edge_index_a, edge_attr_a))<br>
x_b = F.relu(self.conv1(x_b, edge_index_b, edge_attr_b))<br>
x_a, edge_index_a, edge_attr_a, batch = self.pool1(x_a, edge_index_a, edge_attr_a, batch)<br>
x_b, edge_index_b, edge_attr_b, batch = self.pool1(x_b, edge_index_b, edge_attr_b, batch)<br>
x1_a = torch.cat([gmp(x_a, batch), gap(x_a, batch)], dim=1)<br>
x1_b = torch.cat([gmp(x_b, batch), gap(x_b, batch)], dim=1)<br>
…<br>
x = F.relu(self.lin1(torch.cat((x1_a, x1_b), dim=1)))<br>
…<br>
return x</p>
</blockquote>
<p>Actually, I’m not sure that this makes sense - how the gradients and the layers weights will be computed?</p>
<p><strong>b.</strong> Build a new model that duplicate all the conv and pool layers (For example x_a will be passed through conv1_a, pool1_a etc. and x_b through conv1_b, poo1_b etc.).<br>
and concatenate their outputs before the linear layers.</p>
<p>The problem (or not?) with this approach - conv1_a and conv1_b will hold different weights - means that I will have to train the model with some data transformation that switches between subgraphs a and subgraphs b (because I want the model to be invariant for the order of the graphs).<br>
Another problem - I guess that this model will be heavier than the original one.</p>
<p><strong>c</strong> Create a main model that holds two separate original models (without the linear layers) and pass subgraph a through model_a and subgraph b through model_b. Concatenate the results and pass in linear layers i the main model.<br>
This approach is similar to b (same problems) - but I want to make sure that it makes sense.</p>
<p>So which model are you guys recommend?<br>
Is there other good practice to deal with such graphs?</p>
<p>Thanks!</p>