Replace FFNN in GCN-FFNN model with other regressor

[Teslim]

Hi community, thanks for the wonderful work that you are doing. Is it possible to replace the nn.Linear in the code below with other common regressors like RandomForest from Sckit-Learn? Any suggestion on how?

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import dgl
from dgl.nn import GraphConv
from sklearn.ensemble import RandomForestRegressor

class GCN(nn.Module):
    def __init__(self, in_feats, h_feats, num_classes):
        super(GCN, self).__init__()
        if  torch.cuda.is_available():
            self.device = torch.device('cuda:0')
        else:
            self.device = torch.device('cpu')
        self.conv = GraphConv(in_feats, h_feats)
        if torch.cuda.is_available():
            self.conv = self.conv.cuda()
        self.readout1 = nn.Linear(h_feats, h_feats).to(self.device)
        self.readout2 = nn.Linear(h_feats, num_classes).to(self.device)

    def forward(self, g, in_feat):
        with g.local_scope():
            h = self.conv(g, in_feat)
            h = F.relu(h)
            g.ndata['h'] = h
            meanNodes = dgl.mean_nodes(g, 'h')
            output = F.relu(self.readout1(meanNodes))
            output = self.readout2(output)
        return output

[minjie]

Replacing it with sklearn.RandomForestRegressor will break auto-differentiation, so the model will not train properly. Maybe this post on PyTorch forum will be helpful? How can I use KNN, Random Forest models in Pytorch? - #3 by pydv - vision - PyTorch Forums