def test_profile():
root = osp.join('/', 'tmp', str(random.randrange(sys.maxsize)))
dataset = Planetoid(root, 'PubMed')
data = dataset[0].cuda()
model = GraphSAGE(dataset.num_features,
hidden_channels=64,
num_layers=3,
out_channels=dataset.num_classes).cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
@profileit()
def train(model, x, edge_index, y):
model.train()
optimizer.zero_grad()
out = model(x, edge_index)
loss = F.cross_entropy(out, y)
loss.backward()
return float(loss)
@timeit()
@torch.no_grad()
def test(model, x, edge_index, y):
model.eval()
out = model(x, edge_index).argmax(dim=-1)
return int((out == y).sum()) / y.size(0)
stats_list = []
for epoch in range(5):
_, stats = train(model, data.x, data.edge_index, data.y)
assert len(stats) == 6
assert stats.time > 0
assert stats.max_allocated_cuda > 0
assert stats.max_reserved_cuda > 0
assert stats.max_active_cuda > 0
assert stats.nvidia_smi_free_cuda > 0
assert stats.nvidia_smi_used_cuda > 0
_, time = test(model, data.x, data.edge_index, data.y)
assert time > 0
if epoch >= 2: # Warm-up
stats_list.append(stats)
stats_summary = get_stats_summary(stats_list)
assert len(stats_summary) == 7
assert stats_summary.time_mean > 0
assert stats_summary.time_std > 0
assert stats_summary.max_allocated_cuda > 0
assert stats_summary.max_reserved_cuda > 0
assert stats_summary.max_active_cuda > 0
assert stats_summary.min_nvidia_smi_free_cuda > 0
assert stats_summary.max_nvidia_smi_used_cuda > 0
shutil.rmtree(root)
def test_to_hetero_with_basic_model():
x_dict = {
'paper': torch.randn(100, 16),
'author': torch.randn(100, 16),
}
edge_index_dict = {
('paper', 'cites', 'paper'):
torch.randint(100, (2, 200), dtype=torch.long),
('paper', 'written_by', 'author'):
torch.randint(100, (2, 200), dtype=torch.long),
('author', 'writes', 'paper'):
torch.randint(100, (2, 200), dtype=torch.long),
}
metadata = list(x_dict.keys()), list(edge_index_dict.keys())
model = GraphSAGE((-1, -1), 32, num_layers=3)
model = to_hetero(model, metadata, debug=False)
out = model(x_dict, edge_index_dict)
assert isinstance(out, dict) and len(out) == 2
model = GAT((-1, -1), 32, num_layers=3, add_self_loops=False)
model = to_hetero(model, metadata, debug=False)
out = model(x_dict, edge_index_dict)
assert isinstance(out, dict) and len(out) == 2
def __init__(self, in_channels: int, out_channels: int,
hidden_channels: int = 256, num_layers: int = 2,
dropout: float = 0.5):
super().__init__()
self.gnn = GraphSAGE(in_channels, hidden_channels, num_layers,
out_channels, dropout=dropout,
norm=BatchNorm1d(hidden_channels))
self.train_acc = Accuracy()
self.val_acc = Accuracy()
self.test_acc = Accuracy()
def test_torch_profile(get_dataset):
dataset = get_dataset(name='PubMed')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
data = dataset[0].to(device)
model = GraphSAGE(dataset.num_features,
hidden_channels=64,
num_layers=3,
out_channels=dataset.num_classes).to(device)
model.eval()
@timeit()
def inference_e2e(model, data):
model(data.x, data.edge_index)
@torch_profile()
def inference_profile(model, data):
model(data.x, data.edge_index)
for epoch in range(3):
inference_e2e(model, data)
inference_profile(model, data)
rename_profile_file('test_profile')
assert os.path.exists('profile-test_profile.json')
batch_size=2048,
shuffle=True,
neg_sampling_ratio=0.5,
num_neighbors=[10, 10],
num_workers=6,
persistent_workers=True)
# Evaluation loaders (one datapoint corresponds to a graph)
train_loader = DataLoader(train_dataset, batch_size=2)
val_loader = DataLoader(val_dataset, batch_size=2)
test_loader = DataLoader(test_dataset, batch_size=2)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GraphSAGE(
in_channels=train_dataset.num_features,
hidden_channels=64,
num_layers=2,
out_channels=64,
).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.005)
def train():
model.train()
total_loss = total_examples = 0
for data in tqdm.tqdm(loader):
data = data.to(device)
optimizer.zero_grad()
h = model(data.x, data.edge_index)
h_src = h[data.edge_label_index[0]]
def __init__(self, metadata, hidden_channels, num_layers, output_channels):
super().__init__()
self.model = to_hetero(
GraphSAGE((-1, -1), hidden_channels, num_layers, output_channels),
metadata)