dmlc--dgl
RecLusIve-F
0f2ff47de6
* Add BGRL example * Update README.md * Update utils.py * Update * Update utils.py
82 行
2.8 KiB
Python
82 行
2.8 KiB
Python
import copy
|
|
import torch
|
|
import numpy as np
|
|
from dgl.dataloading import GraphDataLoader
|
|
from dgl.transforms import Compose, DropEdge, FeatMask, RowFeatNormalizer
|
|
from dgl.data import CoauthorCSDataset, CoauthorPhysicsDataset, AmazonCoBuyPhotoDataset, AmazonCoBuyComputerDataset, PPIDataset, WikiCSDataset
|
|
|
|
|
|
class CosineDecayScheduler:
|
|
def __init__(self, max_val, warmup_steps, total_steps):
|
|
self.max_val = max_val
|
|
self.warmup_steps = warmup_steps
|
|
self.total_steps = total_steps
|
|
|
|
def get(self, step):
|
|
if step < self.warmup_steps:
|
|
return self.max_val * step / self.warmup_steps
|
|
elif self.warmup_steps <= step <= self.total_steps:
|
|
return self.max_val * (1 + np.cos((step - self.warmup_steps) * np.pi /
|
|
(self.total_steps - self.warmup_steps))) / 2
|
|
else:
|
|
raise ValueError('Step ({}) > total number of steps ({}).'.format(step, self.total_steps))
|
|
|
|
|
|
def get_graph_drop_transform(drop_edge_p, feat_mask_p):
|
|
transforms = list()
|
|
|
|
# make copy of graph
|
|
transforms.append(copy.deepcopy)
|
|
|
|
# drop edges
|
|
if drop_edge_p > 0.:
|
|
transforms.append(DropEdge(drop_edge_p))
|
|
|
|
# drop features
|
|
if feat_mask_p > 0.:
|
|
transforms.append(FeatMask(feat_mask_p, node_feat_names=['feat']))
|
|
|
|
return Compose(transforms)
|
|
|
|
|
|
def get_wiki_cs(transform=RowFeatNormalizer(subtract_min=True)):
|
|
dataset = WikiCSDataset(transform=transform)
|
|
g = dataset[0]
|
|
std, mean = torch.std_mean(g.ndata['feat'], dim=0, unbiased=False)
|
|
g.ndata['feat'] = (g.ndata['feat'] - mean) / std
|
|
|
|
return [g]
|
|
|
|
|
|
def get_ppi():
|
|
train_dataset = PPIDataset(mode='train')
|
|
val_dataset = PPIDataset(mode='valid')
|
|
test_dataset = PPIDataset(mode='test')
|
|
train_val_dataset = [i for i in train_dataset] + [i for i in val_dataset]
|
|
for idx, data in enumerate(train_val_dataset):
|
|
data.ndata['batch'] = torch.zeros(data.number_of_nodes()) + idx
|
|
data.ndata['batch'] = data.ndata['batch'].long()
|
|
|
|
g = list(GraphDataLoader(train_val_dataset, batch_size=22, shuffle=True))
|
|
|
|
return g, PPIDataset(mode='train'), PPIDataset(mode='valid'), test_dataset
|
|
|
|
|
|
def get_dataset(name, transform=RowFeatNormalizer(subtract_min=True)):
|
|
dgl_dataset_dict = {
|
|
'coauthor_cs': CoauthorCSDataset,
|
|
'coauthor_physics': CoauthorPhysicsDataset,
|
|
'amazon_computers': AmazonCoBuyComputerDataset,
|
|
'amazon_photos': AmazonCoBuyPhotoDataset,
|
|
'wiki_cs': get_wiki_cs,
|
|
'ppi': get_ppi
|
|
}
|
|
|
|
dataset_class = dgl_dataset_dict[name]
|
|
train_data, val_data, test_data = None, None, None
|
|
if name != 'ppi':
|
|
dataset = dataset_class(transform=transform)
|
|
else:
|
|
dataset, train_data, val_data, test_data = dataset_class()
|
|
|
|
return dataset, train_data, val_data, test_data |