def __init__(self, in_feats, n_hidden, n_classes):
super().__init__()
self.layers = nn.ModuleList()
self.layers.append(dglnn.SAGEConv(in_feats, n_hidden, 'mean'))
self.layers.append(dglnn.SAGEConv(n_hidden, n_hidden, 'mean'))
self.layers.append(dglnn.SAGEConv(n_hidden, n_classes, 'mean'))
self.dropout = nn.Dropout(0.5)
def __init__(self, in_feats, hid_feats, out_feats):
super().__init__()
self.conv1 = dglnn.SAGEConv(in_feats=in_feats,
out_feats=hid_feats,
aggregator_type="mean")
self.conv2 = dglnn.SAGEConv(in_feats=hid_feats,
out_feats=out_feats,
aggregator_type="mean")
def __init__(self, in_size, hid_size, out_size):
super().__init__()
self.layers = nn.ModuleList()
# three-layer GraphSAGE-mean
self.layers.append(dglnn.SAGEConv(in_size, hid_size, 'mean'))
self.layers.append(dglnn.SAGEConv(hid_size, hid_size, 'mean'))
self.layers.append(dglnn.SAGEConv(hid_size, out_size, 'mean'))
self.dropout = nn.Dropout(0.5)
self.hid_size = hid_size
self.out_size = out_size
def __init__(self, in_feats, hid_feats, out_feats):
super().__init__()
self.conv1 = dglnn.SAGEConv(in_feats=in_feats,
out_feats=hid_feats,
aggregator_type='pool')
self.conv2 = dglnn.SAGEConv(in_feats=hid_feats,
out_feats=1000,
aggregator_type='pool')
self.conv3 = dglnn.SAGEConv(in_feats=1000,
out_feats=out_feats,
aggregator_type='pool')
def __init__(self, in_feats, n_hidden):
super().__init__()
self.n_hidden = n_hidden
self.layers = nn.ModuleList()
self.layers.append(dglnn.SAGEConv(in_feats, n_hidden, 'mean'))
self.layers.append(dglnn.SAGEConv(n_hidden, n_hidden, 'mean'))
self.layers.append(dglnn.SAGEConv(n_hidden, n_hidden, 'mean'))
self.predictor = nn.Sequential(nn.Linear(n_hidden,
n_hidden), nn.ReLU(),
nn.Linear(n_hidden, n_hidden),
nn.ReLU(), nn.Linear(n_hidden, 1))
def init(self, in_feats, n_hidden, n_classes, n_layers, activation,
dropout):
self.n_layers = n_layers
self.n_hidden = n_hidden
self.n_classes = n_classes
self.layers = nn.ModuleList()
self.layers.append(dglnn.SAGEConv(in_feats, n_hidden, 'mean'))
for i in range(1, n_layers - 1):
self.layers.append(dglnn.SAGEConv(n_hidden, n_hidden, 'mean'))
self.layers.append(dglnn.SAGEConv(n_hidden, n_classes, 'mean'))
self.dropout = nn.Dropout(dropout)
self.activation = activation
def __init__(self, in_feats, n_hidden, n_classes, n_layers, activation,
dropout):
super().__init__()
self.n_layers = n_layers
self.n_hidden = n_hidden
self.n_classes = n_classes
self.layers = nn.ModuleList()
self.bns = nn.ModuleList()
self.res_linears = nn.ModuleList()
self.layers.append(
dglnn.SAGEConv(in_feats,
n_hidden,
'mean',
bias=False,
feat_drop=dropout))
self.bns.append(torch.nn.BatchNorm1d(n_hidden))
self.res_linears.append(torch.nn.Linear(in_feats, n_hidden))
for i in range(1, n_layers - 1):
self.layers.append(
dglnn.SAGEConv(n_hidden,
n_hidden,
'mean',
bias=False,
feat_drop=dropout))
self.bns.append(torch.nn.BatchNorm1d(n_hidden))
self.res_linears.append(torch.nn.Identity())
self.layers.append(
dglnn.SAGEConv(n_hidden,
n_hidden,
'mean',
bias=False,
feat_drop=dropout))
self.bns.append(torch.nn.BatchNorm1d(n_hidden))
self.res_linears.append(torch.nn.Identity())
self.mlp = MLP(in_feats + n_hidden * n_layers,
2 * n_classes,
n_classes,
num_layers=2,
bn=True,
end_up_with_fc=True,
act='LeakyReLU')
self.dropout = nn.Dropout(dropout)
self.activation = activation
self.profile = locals()
def test_sage_conv_bi(idtype, g, aggre_type):
g = g.astype(idtype)
dst_dim = 5 if aggre_type != 'gcn' else 10
sage = nn.SAGEConv((10, dst_dim), 2, aggre_type)
feat = (F.randn((g.number_of_src_nodes(), 10)),
F.randn((g.number_of_dst_nodes(), dst_dim)))
init_params = sage.init(jax.random.PRNGKey(2666), g, feat)
h = sage.apply(init_params, g, feat)
assert h.shape[-1] == 2
assert h.shape[0] == g.number_of_dst_nodes()
def __init__(self, src_id, dst_id, in_c, hid_c, n_layers, device):
super(SAGEModel, self).__init__()
self.graph = dgl.graph((src_id, dst_id), device=device)
self.gcn = nn.ModuleList([
gnn.SAGEConv(in_c if i == 0 else hid_c, hid_c, "pool")
for i in range(n_layers)
])
self.residual = nn.ModuleList([
nn.Identity() if i != 0 else nn.Linear(in_c, hid_c)
for i in range(n_layers)
])
def test_sage_conv2(idtype):
# TODO: add test for blocks
# Test the case for graphs without edges
g = dgl.heterograph({('_U', '_E', '_V'): ([], [])}, {'_U': 5, '_V': 3})
g = g.astype(idtype)
sage = nn.SAGEConv((3, 3), 2, 'gcn')
feat = (F.randn((5, 3)), F.randn((3, 3)))
init_params = sage.init(
jax.random.PRNGKey(2666), g,
(F.copy_to(feat[0], F.ctx()), F.copy_to(feat[1], F.ctx())))
h = sage.apply(init_params, g,
(F.copy_to(feat[0], F.ctx()), F.copy_to(feat[1], F.ctx())))
assert h.shape[-1] == 2
assert h.shape[0] == 3
for aggre_type in ['mean', 'pool']:
sage = nn.SAGEConv((3, 1), 2, aggre_type)
feat = (F.randn((5, 3)), F.randn((3, 1)))
init_params = sage.init(jax.random.PRNGKey(2666), g, feat)
h = sage.apply(init_params, g, feat)
assert h.shape[-1] == 2
assert h.shape[0] == 3
def test_sage_conv(idtype, g, aggre_type):
sage = nn.SAGEConv(5, 10, aggre_type)
feat = F.randn((g.number_of_nodes(), 5))
init_params = sage.init(jax.random.PRNGKey(2666), g, feat)
h = sage.apply(init_params, g, feat)
assert h.shape[-1] == 10
