Support for quantize the graph neural networks (GNNs)

[chumingqian]

Hi, there,

May i ask the optimizer is that support for quantize the graph neural networks (GNNs)?

And if not support, why not?

[i-riyad]

Which GNN you are interested in? Modelopt can quantize GNNs, see the example snippet below for SchNet model.
Let us know if you looking for any other specific GNN support.

import torch
from torch_geometric.nn.models import SchNet
from torch_geometric.data import Data
from torch_geometric.loader import DataLoader
import modelopt.torch.quantization as mtq

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

model = SchNet(hidden_channels=128, num_filters=128, num_interactions=6, 
               num_gaussians=50, cutoff=10.0).to(device)

molecules = []
for _ in range(100):
    n_atoms = torch.randint(5, 20, (1,)).item()
    molecules.append(Data(
        z=torch.randint(1, 10, (n_atoms,)),
        pos=torch.randn(n_atoms, 3),
        y=torch.randn(1)
    ))

loader = DataLoader(molecules, batch_size=32, shuffle=False)

def calibrate(m):
    m.eval()
    with torch.no_grad():
        for i, batch in enumerate(loader):
            if i >= 2: break
            batch = batch.to(device)
            _ = m(z=batch.z, pos=batch.pos, batch=batch.batch)

quantized_model = mtq.quantize(model, mtq.INT8_DEFAULT_CFG, calibrate)

[chumingqian]

Thanks for your guide, I used the the GATconv in mine network.

from torch_geometric.nn import GATConv;

class GATLayer(nn.Module):
    def __init__(self, in_channels, out_channels, heads=2):
        super(GATLayer, self).__init__()
        self.gat = GATConv(in_channels, out_channels, heads=heads, concat=False, edge_dim= 32 )#16

    def forward(self, x, edge_index, edge_attr):
        out, attn_weights = self.gat(x, edge_index, edge_attr=edge_attr, return_attention_weights=True)
        return out, attn_weights

class GATModel(nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super(GATModel, self).__init__()
        self.gat1 = GATLayer(in_channels, hidden_channels)
        self.gat2 = GATLayer(hidden_channels, out_channels)

    def forward(self, x, edge_index, edge_attr):
        device = next(self.parameters()).device

        x, edge_index, edge_attr = x.to(device), edge_index.to(device), edge_attr.to(device)

        x_conf, attn_weights1 = self.gat1(x, edge_index, edge_attr)
        x2_in = torch.relu(x_conf)
        x_cls, attn_weights2 = self.gat2(x2_in, edge_index, edge_attr)
        return x_cls,  (attn_weights1, attn_weights2)