The running result of configuring the aggregate function in nn.conv.gin_conv to max cannot be reproduced!

[SagacityPuppy]

When I deploy the GIN model in my project, the results can be reproduced when using the default add aggregator or mean aggregator, but when adjusted to the max aggregator, the results cannot be reproduced and appear random. I'd appreciate it if you could tell me what the problem is. (all the random seeds are fixed)

[SagacityPuppy]

I have already using SparseTensor and not using global_add_pool instead of torch.nn.functional.adaptive_max_pool1d. And all scripts and platform configurations have not changed. Just modifying the aggregator will cause uncertainty. I guess the problem may not be with the pooling operation #8266.

[rusty1s]

Do you mind sharing your model and the input data you are using?

[SagacityPuppy]

Here are my demo scripts, running on the Ubuntu Laptop with the key configuration:

• pytorch 2.0.0 (py3.10_cuda11.8_cudnn8.7.0_0)
• pyg 2.4.0 (py310_torch_2.0.0_cu118)
• pytorch-scatter 2.1.2
• pytorch-cluster 1.6.3
• pytorch-sparse 0.6.18

---

import os
import torch
import numpy as np
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import sklearn.metrics as metrics

from torch import Tensor
from typing import Callable, Union
from torch.nn import Sequential as Seq, Dropout
from torch.nn import  Linear as Lin, BatchNorm1d as BN1d

from torch_geometric.nn.conv import MessagePassing
from torch_geometric.nn.inits import reset
from torch_geometric.utils import spmm
from torch_geometric.typing import Adj, OptPairTensor, Size, SparseTensor
import warnings
warnings.filterwarnings('ignore', '.*Sparse CSR tensor support is in beta state.*')

def seed_all(seed: int=2020, accurate: int=10):
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.cuda.manual_seed(seed)
    torch.set_printoptions(accurate)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    os.environ['PYTHONHASHSEED'] = str(seed)


class GINConv(MessagePassing):  # origin `GINConv` module
    def __init__(self, nn: Callable, eps: float = 0., train_eps: bool = True,
                 **kwargs):
        kwargs.setdefault('aggr', 'mean')  # 'add', 'mean', 'max'
        super().__init__(**kwargs)
        self.nn = nn
        self.initial_eps = eps
        if train_eps:
            self.eps = torch.nn.Parameter(torch.empty(1))
        else:
            self.register_buffer('eps', torch.empty(1))
        self.reset_parameters()

    def reset_parameters(self):
        super().reset_parameters()
        reset(self.nn)
        self.eps.data.fill_(self.initial_eps)


    def forward(
        self,
        x: Union[Tensor, OptPairTensor],
        edge_index: Adj,
        size: Size = None,
    ) -> Tensor:

        if isinstance(x, Tensor):
            x = (x, x)

        # propagate_type: (x: OptPairTensor)
        out = self.propagate(edge_index, x=x, size=size)

        x_r = x[1]
        if x_r is not None:
            out = out + (1 + self.eps) * x_r

        return self.nn(out)


    def message(self, x_i: Tensor, x_j: Tensor) -> Tensor:
        return x_j

    def message_and_aggregate(self, adj_t: SparseTensor,
                              x: OptPairTensor) -> Tensor:
        if isinstance(adj_t, SparseTensor):
            adj_t = adj_t.set_value(None, layout=None)
        return spmm(adj_t, x[0], reduce=self.aggr)

    def __repr__(self) -> str:
        return f'{self.__class__.__name__}(nn={self.nn})'


class ConvLayer(nn.Module):
    def __init__(self,
        in_channels:    int,
        out_channels:   int,
        bias:           bool=False,
        inplace:		bool=False):
        super(ConvLayer, self).__init__()
        self.net = Seq(
                    Lin(in_features     =in_channels,
                        out_features    =out_channels,
                        bias            =bias),  # Input as (L, C)
                    BN1d(out_channels),
                    nn.LeakyReLU(negative_slope=0.2, inplace=inplace)
                    )

    def forward(self, x):
        return self.net(x)


class demo(torch.nn.Module):
    def __init__(self, *args, **kwargs):
        super(demo, self).__init__()
        self.input_dim = args[0]
        self.width = args[1]

        self.ipl = False  # inplace for act_fun
        self.act = nn.LeakyReLU(negative_slope=0.2, inplace=self.ipl)

        self.gnn1 = GINConv(ConvLayer(16, 64, inplace=self.ipl))
        self.gnn2 = GINConv(ConvLayer(64, 128, inplace=self.ipl))
        self.conv = ConvLayer(192, 512, inplace=self.ipl)
        self.head = Seq(Lin(512*2, 512, bias=False), BN1d(512), self.act,
                        Dropout(p=0.5),
                        Lin(512, 256), BN1d(256), self.act,
                        Dropout(p=0.5),
                        Lin(256, 6))

    def forward(self, x, edge_index):
        B, L, _ = x.shape

        x1 = self.act(self.gnn1(x.view(B*L, -1), edge_index))

        x2 = self.act(self.gnn2(x1.view(B*L, -1), edge_index))

        x = torch.cat((x1, x2), dim=-1)

        x = self.conv(x).view(B, L, -1).transpose(-2,-1)
        x1 = F.adaptive_max_pool1d(x, 1).view(B, -1)
        x2 = F.adaptive_avg_pool1d(x, 1).view(B, -1)
        x = torch.cat((x1, x2), 1)

        x = self.head(x)
        return x


if __name__ == '__main__':

    seed_all()
    L, C = 1200, 16
    epoch = 10
    device = torch.device('cuda')
    x = torch.randn(10, L, C)
    y = torch.randint(4, (10,1))
    edge_index = torch.randint(L,(2,7032))
    adj_t = SparseTensor(row=edge_index[1], col=edge_index[0], sparse_sizes=(10*L,10*L))

    model = demo(C,L).to(device)
    opt = optim.Adam(model.parameters(), lr=0.1)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(0, epoch):
        model.train()
        data, label = x.to(device), y.to(device).squeeze()
        opt.zero_grad()
        logits = model(data, adj_t.to(device))
        loss = criterion(logits, label.long())
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 1)
        opt.step()
        preds = logits.max(dim=1)[1]
        loss= loss.item()
        true = label.cpu().numpy()
        pred = preds.detach().cpu().numpy()
        acc = metrics.accuracy_score(true, pred)
        outstr = 'loss: %.6f, acc: %.4f' %(loss, acc)
        print(outstr)

---

The results based on sum, avg and max aggregator are as followed (seed=2020):

• Sum:

  ===> [1st try]
  loss: 2.259676, acc: 0.1000
  loss: 1.806254, acc: 0.5000
  loss: 2.290560, acc: 0.3000
  loss: 1.858529, acc: 0.5000
  loss: 1.350612, acc: 0.6000
  loss: 1.029084, acc: 0.6000
  loss: 1.189798, acc: 0.7000
  loss: 0.609733, acc: 0.8000
  loss: 1.278280, acc: 0.5000
  loss: 2.653089, acc: 0.6000
  ===> [2nd try]
  loss: 2.259676, acc: 0.1000
  loss: 1.806254, acc: 0.5000
  loss: 2.290560, acc: 0.3000
  loss: 1.858529, acc: 0.5000
  loss: 1.350612, acc: 0.6000
  loss: 1.029084, acc: 0.6000
  loss: 1.189798, acc: 0.7000
  loss: 0.609733, acc: 0.8000
  loss: 1.278280, acc: 0.5000
  loss: 2.653089, acc: 0.6000

• Avg:

  ===> [1st try]
  loss: 2.259676, acc: 0.1000
  loss: 1.806254, acc: 0.5000
  loss: 2.290560, acc: 0.3000
  loss: 1.858529, acc: 0.5000
  loss: 1.350612, acc: 0.6000
  loss: 1.029084, acc: 0.6000
  loss: 1.189798, acc: 0.7000
  loss: 0.609733, acc: 0.8000
  loss: 1.278280, acc: 0.5000
  loss: 2.653089, acc: 0.6000
  ===> [2nd try]
  loss: 2.259676, acc: 0.1000
  loss: 1.806254, acc: 0.5000
  loss: 2.290560, acc: 0.3000
  loss: 1.858529, acc: 0.5000
  loss: 1.350612, acc: 0.6000
  loss: 1.029084, acc: 0.6000
  loss: 1.189798, acc: 0.7000
  loss: 0.609733, acc: 0.8000
  loss: 1.278280, acc: 0.5000
  loss: 2.653089, acc: 0.6000

• Max:

  ===> [1st try]
  loss: 2.353178, acc: 0.2000
  loss: 1.550491, acc: 0.4000
  loss: 2.354827, acc: 0.3000
  loss: 1.416273, acc: 0.6000
  loss: 2.028093, acc: 0.4000
  loss: 1.838474, acc: 0.5000
  loss: 0.326978, acc: 0.8000
  loss: 0.438439, acc: 0.8000
  loss: 0.426627, acc: 0.9000
  loss: 0.348055, acc: 0.9000
  ===> [2nd try]
  loss: 2.353178, acc: 0.2000
  loss: 1.550490, acc: 0.4000
  loss: 2.354824, acc: 0.3000
  loss: 1.416710, acc: 0.6000
  loss: 2.018174, acc: 0.4000
  loss: 1.818063, acc: 0.5000
  loss: 0.408936, acc: 0.7000
  loss: 0.343728, acc: 0.8000
  loss: 0.244022, acc: 0.9000
  loss: 1.427194, acc: 0.8000

---

Hope your answers and many thanks!!

[rusty1s]

Thank you. I can reproduce this. The issue is that segment_max is not deterministic in the backward pass if the maximum is shared across neighbors (in which case it picks one element randomly). This is not sadly not super trivial to fix.

[SagacityPuppy]

Thank you very much for your answers.

Please forgive mine extra asking:

• Does this mean that the implementation of segment_max is different from the typical max-pooling in PyTorch?
• Even if I fix all the random seeds, will the program still pick random elements in the backward pass?
• In addition, segment_min should also have the same problem?

Looking forward to your insights and thanks in advance ~