dmlc--dgl
Mufei Li
828a5e5bc6
* First commit * Update * Update splitters * Update * Update * Update * Update * Update * Update * Migrate ACNN * Fix * Fix * Update * Update * Update * Update * Update * Update * Finish classification * Update * Fix * Update * Update * Update * Fix * Fix * Fix * Update * Update * Update * trigger CI * Fix CI * Update * Update * Update * Add default values * Rename * Update deprecation message
86 行
3.0 KiB
Python
86 行
3.0 KiB
Python
import dgl
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from dgl import DGLGraph
|
|
from dgllife.model.readout import *
|
|
|
|
def test_graph1():
|
|
"""Graph with node features"""
|
|
g = DGLGraph([(0, 1), (0, 2), (1, 2)])
|
|
return g, torch.arange(g.number_of_nodes()).float().reshape(-1, 1)
|
|
|
|
def test_graph2():
|
|
"Batched graph with node features"
|
|
g1 = DGLGraph([(0, 1), (0, 2), (1, 2)])
|
|
g2 = DGLGraph([(0, 1), (1, 2), (1, 3), (1, 4)])
|
|
bg = dgl.batch([g1, g2])
|
|
return bg, torch.arange(bg.number_of_nodes()).float().reshape(-1, 1)
|
|
|
|
def test_weighted_sum_and_max():
|
|
if torch.cuda.is_available():
|
|
device = torch.device('cuda:0')
|
|
else:
|
|
device = torch.device('cpu')
|
|
|
|
g, node_feats = test_graph1()
|
|
g, node_feats = g.to(device), node_feats.to(device)
|
|
bg, batch_node_feats = test_graph2()
|
|
bg, batch_node_feats = bg.to(device), batch_node_feats.to(device)
|
|
model = WeightedSumAndMax(in_feats=1).to(device)
|
|
assert model(g, node_feats).shape == torch.Size([1, 2])
|
|
assert model(bg, batch_node_feats).shape == torch.Size([2, 2])
|
|
|
|
def test_attentive_fp_readout():
|
|
if torch.cuda.is_available():
|
|
device = torch.device('cuda:0')
|
|
else:
|
|
device = torch.device('cpu')
|
|
|
|
g, node_feats = test_graph1()
|
|
g, node_feats = g.to(device), node_feats.to(device)
|
|
bg, batch_node_feats = test_graph2()
|
|
bg, batch_node_feats = bg.to(device), batch_node_feats.to(device)
|
|
model = AttentiveFPReadout(feat_size=1,
|
|
num_timesteps=1).to(device)
|
|
assert model(g, node_feats).shape == torch.Size([1, 1])
|
|
assert model(bg, batch_node_feats).shape == torch.Size([2, 1])
|
|
|
|
def test_mlp_readout():
|
|
if torch.cuda.is_available():
|
|
device = torch.device('cuda:0')
|
|
else:
|
|
device = torch.device('cpu')
|
|
|
|
g, node_feats = test_graph1()
|
|
g, node_feats = g.to(device), node_feats.to(device)
|
|
bg, batch_node_feats = test_graph2()
|
|
bg, batch_node_feats = bg.to(device), batch_node_feats.to(device)
|
|
|
|
model = MLPNodeReadout(node_feats=1,
|
|
hidden_feats=2,
|
|
graph_feats=3,
|
|
activation=F.relu,
|
|
mode='sum').to(device)
|
|
assert model(g, node_feats).shape == torch.Size([1, 3])
|
|
assert model(bg, batch_node_feats).shape == torch.Size([2, 3])
|
|
|
|
model = MLPNodeReadout(node_feats=1,
|
|
hidden_feats=2,
|
|
graph_feats=3,
|
|
mode='max').to(device)
|
|
assert model(g, node_feats).shape == torch.Size([1, 3])
|
|
assert model(bg, batch_node_feats).shape == torch.Size([2, 3])
|
|
|
|
model = MLPNodeReadout(node_feats=1,
|
|
hidden_feats=2,
|
|
graph_feats=3,
|
|
mode='mean').to(device)
|
|
assert model(g, node_feats).shape == torch.Size([1, 3])
|
|
assert model(bg, batch_node_feats).shape == torch.Size([2, 3])
|
|
|
|
if __name__ == '__main__':
|
|
test_weighted_sum_and_max()
|
|
test_attentive_fp_readout()
|
|
test_mlp_readout()
|