DGL&PyTorch in C++

[Kenny]

Hi all,

I’m currently training my model using DGL and PyTorch. Can I pack my model for C++ inference?
I know there’s a TorchScript tool that can pack models, but I customized a single layer using DGL to pass through the message. Can this also be packed by TorchScript?
If it can, the input of my model is a DGL graph and some labels. How can I handle the DGL graph structure in C++?

Thanks for reading this question.

[frozenbugs]

How did you implement your DGL layer? If you use Tutorials: dgl.sparse — DGL 2.2.1 documentation, I think you can do it follow Torch’s recipe. For MFG based implementation, we don’t support it yet. Some migration need to be done to support it.

[Kenny]

Below is how I implement my module.

class GCN(nn.Module):
def init(self, in_feats, h_feats):
super(GCN, self).init()
self.transform_linear = nn.Linear(in_feats + 16 + 7, h_feats, bias=True)

def forward(self, g, h):
    g = g.to(device)
    g.ndata['h'] = torch.zeros(g.number_of_nodes(), 128, device=device)
    g.ndata['reduce'] = torch.zeros(g.number_of_nodes(), 128 + 16, device=device)
    g.edata['msg'] = torch.zeros(g.number_of_edges(), 144, device=device)
    topological_order = list(dgl.topological_nodes_generator(g))
    for order in topological_order:
        for node_id in order:
            node_id = node_id.to(device) 
            
            g.ndata['h'][node_id] = h[node_id]
            
            src, _, in_edges = g.in_edges(node_id, form='all')
            _, dst, out_edges  = g.out_edges(node_id, form='all')

            if len(out_edges) > 0:
                for out_edges_id in out_edges.tolist():
                    controllability = g.edata['edge_features'][out_edges_id, 0]
                    node_data = g.ndata['h'][node_id].unsqueeze(0)
                    edge_data = g.edata['edge_trans_cor'][out_edges_id].unsqueeze(0)
                    msg = torch.cat([node_data, edge_data], dim=1)                    
                    msg *= controllability
                    g.edata['msg'][out_edges_id] = msg
            
            if len(in_edges) > 0:
                msg = g.edata['msg'][in_edges]
                reduce_msg = msg.sum(dim=0)
                g.ndata['reduce'][node_id] = reduce_msg
                        
            else:
                edge_data_zero = torch.zeros(16, device=device).unsqueeze(0)
                hidden_feature = g.ndata['h'][node_id].unsqueeze(0)  # 2D
                g.ndata['reduce'][node_id] = torch.cat([hidden_feature, edge_data_zero], dim=1)  

            node_feature = g.ndata['node_features'][node_id].unsqueeze(0)
            reduce_feature = g.ndata['reduce'][node_id].unsqueeze(0)
            transform = torch.cat([reduce_feature, node_feature], dim=1)
            h[node_id] = self.transform_linear(transform)

    return h

class Mymodule(nn.Module):
def init(self, in_feats, h_feats, num_classes):
super(Mymodule, self).init()
self.fc1 = nn.Linear(in_feats, h_feats)
self.conv1 = GCN(h_feats, h_feats)
self.fc2 = nn.Linear(h_feats, 32)
self.fc3 = nn.Linear(32, num_classes)

def forward(self, g, labels):
    g = g.to(device)
    labels = labels.to(device)
    h = self.fc1(labels)
    h = F.leaky_relu(h, negative_slope=0.01)
    h = self.conv1(g, h)
    h = F.leaky_relu(h, negative_slope=0.01)
    h = self.fc2(h)
    h = F.leaky_relu(h, negative_slope=0.01)
    h = self.fc3(h)
    h = F.softmax(h, dim=1)
    return h

model = Mymodule(in_feats, h_feats, num_classes).to(device)
outputs = model(dgl_G, label)

And this is how I declare model and use it, I want to ask:

1. Can I create dgl graph for inference in C++ ?
2. Can I pack this kind of model in torchscript?

Thanks for reading this.