diff --git a/examples/diffpool/README.md b/examples/diffpool/README.md new file mode 100644 index 000000000..e3f41ebdc --- /dev/null +++ b/examples/diffpool/README.md @@ -0,0 +1,27 @@ +# Differentiable Pooling + +- Paper link: [https://arxiv.org/abs/1806.08804]( https://arxiv.org/abs/1806.08804) +- Author's code repo: [https://github.com/RexYing/diffpool](https://github.com/RexYing/diffpool). Note that the original code is + implemented with Pytorch for the paper. + +# Dataset Statics + +| Dataset | # Graphs | # Edges | # Avg of graph size | +|---------|----------|---------|---------------------| +| syn1v2 | 1000 | 181067 | 49.23 | +| D&D | 1168 | 789819 | 268.70 | +| ENZYMES | 600 | 37282 | 32.46 | + +Results +------- + +```bash +python diffpool_trainer.py --dataset syn1v2 --num_pool 1 --epochs 50 +python diffpool_trainer.py --bmname DD +``` + +| Dataset | Paper | Our(pytorch) | Our(tf) | +|---------|-------|--------------|---------| +| syn1v2 | - | 0.80 | | +| D&D | 81.15 | 79.31 | | +| ENZYMES | 64.23 | | | \ No newline at end of file diff --git a/examples/diffpool/diffpool_trainer.py b/examples/diffpool/diffpool_trainer.py new file mode 100644 index 000000000..0e1d8fd0d --- /dev/null +++ b/examples/diffpool/diffpool_trainer.py @@ -0,0 +1,412 @@ +# !/usr/bin/env python +# -*- encoding: utf-8 -*- +""" +@File : diffpool_trainer.py +@Time : 2022/04/05 18:00:00 +@Author : yang yuxin +""" + +import os +os.environ['TL_BACKEND'] = 'torch' # set your backend here, default `tensorflow` +# os.environ["CUDA_VISIBLE_DEVICES"] = "8" +# os.environ['CUDA_LAUNCH_BLOCKING'] = "1" +import sys +sys.path.insert(0, os.path.abspath('../../')) +import pickle +import shutil +import random +import argparse +import tensorlayerx as tlx +import numpy as np +import networkx as nx +import abc +from tensorboardX import SummaryWriter +from tensorlayerx import Variable +import sklearn.metrics as metrics +from gammagl.datasets.dd import prepare_val_data,read_graphfile,GraphSampler +from gammagl.models.diffpool import GcnEncoderGraph + +def evaluate(dataset, model, args, name='Validation', max_num_examples=None): + model.eval() + labels = [] + preds = [] + for batch_idx, data in enumerate(dataset): + adj = Variable(data['adj'].float()) + h0 = Variable(data['feats'].float()) + labels.append(data['label'].long().numpy()) + batch_num_nodes = data['num_nodes'].int().numpy() + assign_input = Variable(data['assign_feats'].float()) + ypred = model(h0, adj, batch_num_nodes, assign_x=assign_input) + indices = tlx.argmax(ypred, axis=1) + preds.append(indices.cpu().data.numpy()) + if max_num_examples is not None: + if (batch_idx + 1) * args.batch_size > max_num_examples: + break + + labels = np.hstack(labels) + preds = np.hstack(preds) + result = { 'acc': metrics.accuracy_score(labels, preds)} + print(name, " accuracy:", result['acc']) + return result + +def gen_prefix(args): + if args.bmname is not None: + name = args.bmname + else: + name = args.dataset + name += '_' + args.method + + name += '_l' + str(args.num_gc_layers) + name += '_h' + str(args.hidden_dim) + '_o' + str(args.output_dim) + if not args.bias: + name += '_nobias' + if len(args.name_suffix) > 0: + name += '_' + args.name_suffix + return name + + +def train(dataset, model, args, same_feat=True, val_dataset=None, test_dataset=None, writer=None, + mask_nodes=True): + writer_batch_idx = [0, 3, 6, 9] + + optimizer = tlx.optimizers.Adam(lr=0.001) + iter = 0 + train_accs = [] + train_epochs = [] + test_accs = [] + val_accs = [] + + for epoch in range(args.num_epochs): + total_time = 0 + avg_loss = 0.0 + model.set_train() + for batch_idx, data in enumerate(dataset): + model.zero_grad() + adj = Variable(data['adj'].float(),"adj",False) + h0 = Variable(data['feats'].float(),"h0",False) + label = Variable(data['label'].long(),"label",False) + batch_num_nodes = data['num_nodes'].int().numpy() if mask_nodes else None + assign_input = Variable(data['assign_feats'].float(),"assign_input",False) + + ypred = model(h0, adj, batch_num_nodes, assign_x=assign_input) + if not args.method == 'soft-assign' or not args.linkpred: + loss = model.loss(ypred, label) + else: + loss = model.loss(ypred, label, adj, batch_num_nodes) + #loss.backward() + + g=optimizer.gradient(loss, weights=model.trainable_weights) + optimizer.apply_gradients(zip(g, model.trainable_weights)) + + iter += 1 + avg_loss += loss + + avg_loss /= batch_idx + 1 + if writer is not None: + writer.add_scalar('loss/avg_loss', avg_loss, epoch) + if args.linkpred: + writer.add_scalar('loss/linkpred_loss', model.link_loss, epoch) + print('Avg loss: ', avg_loss) + result = evaluate(dataset, model, args, name='Train', max_num_examples=100) + train_accs.append(result['acc']) + train_epochs.append(epoch) + if val_dataset is not None: + val_result = evaluate(val_dataset, model, args, name='Validation') + val_accs.append(val_result['acc']) + if test_dataset is not None: + test_result = evaluate(test_dataset, model, args, name='Test') + test_result['epoch'] = epoch + if writer is not None: + writer.add_scalar('acc/train_acc', result['acc'], epoch) + writer.add_scalar('acc/val_acc', val_result['acc'], epoch) + if test_dataset is not None: + writer.add_scalar('acc/test_acc', test_result['acc'], epoch) + return model, val_accs + + +def prepare_data(graphs, args, test_graphs=None, max_nodes=0): + random.shuffle(graphs) + if test_graphs is None: + train_idx = int(len(graphs) * args.train_ratio) + test_idx = int(len(graphs) * (1 - args.test_ratio)) + train_graphs = graphs[:train_idx] + val_graphs = graphs[train_idx: test_idx] + test_graphs = graphs[test_idx:] + else: + train_idx = int(len(graphs) * args.train_ratio) + train_graphs = graphs[:train_idx] + val_graphs = graphs[train_idx:] + + # minibatch + dataset_sampler = GraphSampler(train_graphs, normalize=False, max_num_nodes=max_nodes, + features=args.feature_type) + train_dataset_loader = tlx.dataflow.DataLoader( + dataset_sampler, + batch_size=args.batch_size, + shuffle=True, + num_workers=args.num_workers) + + dataset_sampler = GraphSampler(val_graphs, normalize=False, max_num_nodes=max_nodes, + features=args.feature_type) + val_dataset_loader = tlx.dataflow.DataLoader( + dataset_sampler, + batch_size=args.batch_size, + shuffle=False, + num_workers=args.num_workers) + + dataset_sampler = GraphSampler(test_graphs, normalize=False, max_num_nodes=max_nodes, + features=args.feature_type) + test_dataset_loader = tlx.dataflow.DataLoader( + dataset_sampler, + batch_size=args.batch_size, + shuffle=False, + num_workers=args.num_workers) + + return train_dataset_loader, val_dataset_loader, test_dataset_loader, \ + dataset_sampler.max_num_nodes, dataset_sampler.feat_dim, dataset_sampler.assign_feat_dim + + +class FeatureGen(metaclass=abc.ABCMeta): + @abc.abstractmethod + def gen_node_features(self, G): + pass + +class ConstFeatureGen(FeatureGen): + def __init__(self, val): + self.val = val + + def gen_node_features(self, G): + feat_dict = {i:{'feat': self.val} for i in G.nodes()} + nx.set_node_attributes(G, feat_dict) + +class GaussianFeatureGen(FeatureGen): + def __init__(self, mu, sigma): + self.mu = mu + self.sigma = sigma + + def gen_node_features(self, G): + feat = np.random.multivariate_normal(self.mu, self.sigma, G.number_of_nodes()) + feat_dict = {i:{'feat': feat[i]} for i in range(feat.shape[0])} + nx.set_node_attributes(G, feat_dict) + +def gen_ba(n_range, m_range, num_graphs, feature_generator=None): + graphs = [] + for i in np.random.choice(n_range, num_graphs): + for j in np.random.choice(m_range, 1): + graphs.append(nx.barabasi_albert_graph(i,j)) + + if feature_generator is None: + feature_generator = ConstFeatureGen(0) + for G in graphs: + feature_generator.gen_node_features(G) + return graphs + +def gen_er(n_range, p, num_graphs, feature_generator=None): + graphs = [] + for i in np.random.choice(n_range, num_graphs): + graphs.append(nx.erdos_renyi_graph(i,p)) + + if feature_generator is None: + feature_generator = ConstFeatureGen(0) + for G in graphs: + feature_generator.gen_node_features(G) + return graphs + +def gen_2community_ba(n_range, m_range, num_graphs, inter_prob, feature_generators): + ''' Each community is a BA graph. + Args: + inter_prob: probability of one node connecting to any node in the other community. + ''' + + if feature_generators is None: + mu0 = np.zeros(10) + mu1 = np.ones(10) + sigma0 = np.ones(10, 10) * 0.1 + sigma1 = np.ones(10, 10) * 0.1 + fg0 = GaussianFeatureGen(mu0, sigma0) + fg1 = GaussianFeatureGen(mu1, sigma1) + else: + fg0 = feature_generators[0] + fg1 = feature_generators[1] if len(feature_generators) > 1 else feature_generators[0] + + graphs1 = [] + graphs2 = [] + graphs0 = gen_ba(n_range, m_range, num_graphs, fg0) + graphs1 = gen_ba(n_range, m_range, num_graphs, fg1) + graphs = [] + for i in range(num_graphs): + G = nx.disjoint_union(graphs0[i], graphs1[i]) + n0 = graphs0[i].number_of_nodes() + for j in range(n0): + if np.random.rand() < inter_prob: + target = np.random.choice(G.number_of_nodes() - n0) + n0 + G.add_edge(j, target) + graphs.append(G) + return graphs + +def syn_community1v2(args, writer=None, export_graphs=False): + graphs1 = gen_ba(range(40, 60), range(4, 5), 500, + ConstFeatureGen(np.ones(args.input_dim, dtype=float))) + for G in graphs1: + G.graph['label'] = 0 + + graphs2 = gen_2community_ba(range(20, 30), range(4, 5), 500, 0.3, + [ConstFeatureGen(np.ones(args.input_dim, dtype=float))]) + for G in graphs2: + G.graph['label'] = 1 + + graphs = graphs1 + graphs2 + + train_dataset, val_dataset, test_dataset, max_num_nodes, input_dim, assign_input_dim = prepare_data(graphs, args) + if args.method == 'base': + model = GcnEncoderGraph(input_dim, args.hidden_dim, args.output_dim, 2, + args.num_gc_layers, bn=args.bn) + + train(train_dataset, model, args, val_dataset=val_dataset, test_dataset=test_dataset, + writer=writer) + +def benchmark_task(args, writer=None, feat='node-label'): + graphs = read_graphfile(args.datadir, args.bmname, max_nodes=args.max_nodes) + + if feat == 'node-feat' and 'feat_dim' in graphs[0].graph: + input_dim = graphs[0].graph['feat_dim'] + elif feat == 'node-label' and 'label' in graphs[0].node[0]: + for G in graphs: + for u in G.nodes(): + G.node[u]['feat'] = np.array(G.node[u]['label']) + else: + featgen_const = ConstFeatureGen(np.ones(args.input_dim, dtype=float)) + for G in graphs: + featgen_const.gen_node_features(G) + + train_dataset, val_dataset, test_dataset, max_num_nodes, input_dim, assign_input_dim = \ + prepare_data(graphs, args, max_nodes=args.max_nodes) + if args.method == 'base': + model = GcnEncoderGraph( + input_dim, args.hidden_dim, args.output_dim, args.num_classes, + args.num_gc_layers, bn=args.bn, dropout=args.dropout, args=args).cuda() + + train(train_dataset, model, args, val_dataset=val_dataset, test_dataset=test_dataset, + writer=writer) + evaluate(test_dataset, model, args, 'Validation') + +def node_iter(G): + return G.nodes + +def node_dict(G): + node_dict = G.nodes + return node_dict + +def exp_moving_avg(x, decay=0.9): + shadow = x[0] + a = [shadow] + for v in x[1:]: + shadow -= (1-decay) * (shadow-v) + a.append(shadow) + return a + +def main(prog_args): + import multiprocessing as mp + mp.set_start_method('spawn') + # export scalar data to JSON for external processing + path = os.path.join(prog_args.logdir, gen_prefix(prog_args)) + if os.path.isdir(path): + shutil.rmtree(path) + writer = SummaryWriter(path) + writer = None + if prog_args.bmname is not None: + benchmark_task_val(prog_args, writer=writer) + elif prog_args.dataset is not None: + if prog_args.dataset == 'syn1v2': + syn_community1v2(prog_args, writer=writer) + writer.close() + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='GraphPool arguments.') + io_parser = parser.add_mutually_exclusive_group(required=False) + io_parser.add_argument('--dataset', default='syn1v2',dest='dataset', + help='Input dataset.') + benchmark_parser = io_parser.add_argument_group() + benchmark_parser.add_argument('--bmname', dest='bmname', + help='Name of the benchmark dataset') + softpool_parser = parser.add_argument_group() + softpool_parser.add_argument('--assign-ratio',default=0.1, dest='assign_ratio', type=float, + help='ratio of number of nodes in consecutive layers') + softpool_parser.add_argument('--num-pool', default=1,dest='num_pool', type=int, + help='number of pooling layers') + parser.add_argument('--linkpred', dest='linkpred', action='store_const', + const=True, default=False, + help='Whether link prediction side objective is used') + parser.add_argument('--datadir', dest='datadir', default='data', + help='Directory where benchmark is located') + parser.add_argument('--logdir', dest='logdir', default='log', + help='Tensorboard log directory') + parser.add_argument('--cuda', dest='cuda', default='8', + help='CUDA.') + parser.add_argument('--max-nodes', dest='max_nodes', type=int,default=1000, + help='Maximum number of nodes (ignore graghs with nodes exceeding the number.') + parser.add_argument('--lr', dest='lr', type=float, default=0.001, + help='Learning rate.') + parser.add_argument('--clip', dest='clip', type=float, default=2.0, + help='Gradient clipping.') + parser.add_argument('--batch-size', dest='batch_size', type=int, default=20, + help='Batch size.') + parser.add_argument('--epochs', dest='num_epochs', type=int, default=1000, + help='Number of epochs to train.') + parser.add_argument('--train-ratio', dest='train_ratio', type=float,default=0.8, + help='Ratio of number of graphs training set to all graphs.') + parser.add_argument('--num_workers', dest='num_workers', type=int,default=1, + help='Number of workers to load data.') + parser.add_argument('--feature', dest='feature_type',default='default', + help='Feature used for encoder. Can be: id, deg') + parser.add_argument('--input-dim', dest='input_dim', type=int, default=10, + help='Input feature dimension') + parser.add_argument('--hidden-dim', dest='hidden_dim', type=int, default=20, + help='Hidden dimension') + parser.add_argument('--output-dim', dest='output_dim', type=int,default=20, + help='Output dimension') + parser.add_argument('--num-classes', dest='num_classes', type=int,default=2, + help='Number of label classes') + parser.add_argument('--num-gc-layers', dest='num_gc_layers', type=int,default=3, + help='Number of graph convolution layers before each pooling') + parser.add_argument('--nobn', dest='bn', action='store_const', + const=False, default=True, + help='Whether batch normalization is used') + parser.add_argument('--dropout', dest='dropout', type=float, default=0.0, + help='Dropout rate.') + parser.add_argument('--nobias', dest='bias', action='store_const', + const=False, default=True, + help='Whether to add bias. Default to True.') + parser.add_argument('--no-log-graph', dest='log_graph', action='store_const', + const=False, default=True, + help='Whether disable log graph') + parser.add_argument('--method', dest='method',default='base', + help='Method. Possible values: base') + parser.add_argument('--name-suffix', dest='name_suffix',default='', + help='suffix added to the output filename') + parser.set_defaults(datadir='data', + logdir='log', + dataset='syn1v2', + max_nodes=1000, + cuda='8', + feature_type='default', + lr=0.001, + clip=2.0, + batch_size=20, + num_epochs=50, + train_ratio=0.8, + test_ratio=0.1, + num_workers=1, + input_dim=10, + hidden_dim=20, + output_dim=20, + num_classes=2, + num_gc_layers=3, + dropout=0.0, + method='base', + name_suffix='', + assign_ratio=0.1, + num_pool=1 + ) + prog_args = parser.parse_args() + main(prog_args) diff --git a/gammagl/datasets/dd.py b/gammagl/datasets/dd.py new file mode 100644 index 000000000..94b3e06e7 --- /dev/null +++ b/gammagl/datasets/dd.py @@ -0,0 +1,266 @@ +import networkx as nx +import numpy as np +import os +import re +import tensorlayerx as tlx +import random + + +def node_iter(G): + return G.nodes + +def node_dict(G): + node_dict = G.nodes + return node_dict + +def exp_moving_avg(x, decay=0.9): + shadow = x[0] + a = [shadow] + for v in x[1:]: + shadow -= (1-decay) * (shadow-v) + a.append(shadow) + return a + +class GraphSampler(tlx.dataflow.Dataset): + ''' Sample graphs and nodes in graph + ''' + + def __init__(self, G_list, features='default', normalize=True, assign_feat='default', max_num_nodes=0): + self.adj_all = [] + self.len_all = [] + self.feature_all = [] + self.label_all = [] + + self.assign_feat_all = [] + + if max_num_nodes == 0: + self.max_num_nodes = max([G.number_of_nodes() for G in G_list]) + else: + self.max_num_nodes = max_num_nodes # 1000 + + if features == 'default': + self.feat_dim = node_dict(G_list[0])[0]['feat'].shape[0] # 89 + + # G_list 1052 + for G in G_list: + adj = np.array(nx.to_numpy_matrix(G)) # 159，159 + if normalize: + sqrt_deg = np.diag(1.0 / np.sqrt(np.sum(adj, axis=0, dtype=float).squeeze())) + adj = np.matmul(np.matmul(sqrt_deg, adj), sqrt_deg) + self.adj_all.append(adj) + self.len_all.append(G.number_of_nodes()) + self.label_all.append(G.graph['label']) + # feat matrix: max_num_nodes x feat_dim + if features == 'default': + f = np.zeros((self.max_num_nodes, self.feat_dim), dtype=float) # 1000，89 + for i, u in enumerate(G.nodes()): # 314 + f[i, :] = node_dict(G)[u]['feat'] + self.feature_all.append(f) # 1000，89 + elif features == 'id': + self.feature_all.append(np.identity(self.max_num_nodes)) + elif features == 'deg-num': + degs = np.sum(np.array(adj), 1) + degs = np.expand_dims(np.pad(degs, [0, self.max_num_nodes - G.number_of_nodes()], 0), + axis=1) + self.feature_all.append(degs) + elif features == 'deg': + self.max_deg = 10 + degs = np.sum(np.array(adj), 1).astype(int) + degs[degs > self.max_deg] = self.max_deg + feat = np.zeros((len(degs), self.max_deg + 1)) + feat[np.arange(len(degs)), degs] = 1 + feat = np.pad(feat, ((0, self.max_num_nodes - G.number_of_nodes()), (0, 0)), + 'constant', constant_values=0) + + f = np.zeros((self.max_num_nodes, self.feat_dim), dtype=float) + for i, u in enumerate(node_iter(G)): + f[i, :] = node_dict(G)[u]['feat'] + + feat = np.concatenate((feat, f), axis=1) + + self.feature_all.append(feat) + elif features == 'struct': + self.max_deg = 10 + degs = np.sum(np.array(adj), 1).astype(int) + degs[degs > 10] = 10 + feat = np.zeros((len(degs), self.max_deg + 1)) + feat[np.arange(len(degs)), degs] = 1 + degs = np.pad(feat, ((0, self.max_num_nodes - G.number_of_nodes()), (0, 0)), + 'constant', constant_values=0) + + clusterings = np.array(list(nx.clustering(G).values())) + clusterings = np.expand_dims(np.pad(clusterings, + [0, self.max_num_nodes - G.number_of_nodes()], + 'constant'), + axis=1) + g_feat = np.hstack([degs, clusterings]) + if 'feat' in node_dict(G)[0]: + node_feats = np.array([node_dict(G)[i]['feat'] for i in range(G.number_of_nodes())]) + node_feats = np.pad(node_feats, ((0, self.max_num_nodes - G.number_of_nodes()), (0, 0)), + 'constant') + g_feat = np.hstack([g_feat, node_feats]) + + self.feature_all.append(g_feat) + + if assign_feat == 'id': + self.assign_feat_all.append( + np.hstack((np.identity(self.max_num_nodes), self.feature_all[-1]))) + else: + self.assign_feat_all.append(self.feature_all[-1]) + + self.feat_dim = self.feature_all[0].shape[1] + self.assign_feat_dim = self.assign_feat_all[0].shape[1] + + def __len__(self): + return len(self.adj_all) + + def __getitem__(self, idx): + adj = self.adj_all[idx] + num_nodes = adj.shape[0] + adj_padded = np.zeros((self.max_num_nodes, self.max_num_nodes)) # 1000,1000 + adj_padded[:num_nodes, :num_nodes] = adj + + # use all nodes for aggregation (baseline) + + return {'adj': adj_padded, + 'feats': self.feature_all[idx].copy(), + 'label': self.label_all[idx], + 'num_nodes': num_nodes, + 'assign_feats': self.assign_feat_all[idx].copy()} + + +def read_graphfile(datadir, dataname, max_nodes=None): + ''' Read data from https://ls11-www.cs.tu-dortmund.de/staff/morris/graphkerneldatasets + graph index starts with 1 in file + + Returns: + List of networkx objects with graph and node labels + ''' + prefix = os.path.join(datadir, dataname, dataname) + filename_graph_indic = prefix + '_graph_indicator.txt' + # index of graphs that a given node belongs to + graph_indic = {} + with open(filename_graph_indic) as f: + i = 1 + for line in f: + line = line.strip("\n") + graph_indic[i] = int(line) + i += 1 + + filename_nodes = prefix + '_node_labels.txt' + node_labels = [] + try: + with open(filename_nodes) as f: + for line in f: + line = line.strip("\n") + node_labels += [int(line) - 1] + num_unique_node_labels = max(node_labels) + 1 + except IOError: + print('No node labels') + + filename_node_attrs = prefix + '_node_attributes.txt' + node_attrs = [] + try: + with open(filename_node_attrs) as f: + for line in f: + line = line.strip("\s\n") + attrs = [float(attr) for attr in re.split("[,\s]+", line) if not attr == ''] + node_attrs.append(np.array(attrs)) + except IOError: + print('No node attributes') + + label_has_zero = False + filename_graphs = prefix + '_graph_labels.txt' + graph_labels = [] + + # assume that all graph labels appear in the dataset + # (set of labels don't have to be consecutive) + label_vals = [] + with open(filename_graphs) as f: + for line in f: + line = line.strip("\n") + val = int(line) + # if val == 0: + # label_has_zero = True + if val not in label_vals: + label_vals.append(val) + graph_labels.append(val) + # graph_labels = np.array(graph_labels) + label_map_to_int = {val: i for i, val in enumerate(label_vals)} + graph_labels = np.array([label_map_to_int[l] for l in graph_labels]) + # if label_has_zero: + # graph_labels += 1 + + filename_adj = prefix + '_A.txt' + adj_list = {i: [] for i in range(1, len(graph_labels) + 1)} + index_graph = {i: [] for i in range(1, len(graph_labels) + 1)} + num_edges = 0 + with open(filename_adj) as f: + for line in f: + line = line.strip("\n").split(",") + e0, e1 = (int(line[0].strip(" ")), int(line[1].strip(" "))) + adj_list[graph_indic[e0]].append((e0, e1)) + index_graph[graph_indic[e0]] += [e0, e1] + num_edges += 1 + for k in index_graph.keys(): + index_graph[k] = [u - 1 for u in set(index_graph[k])] + + graphs = [] + for i in range(1, 1 + len(adj_list)): + # indexed from 1 here + G = nx.from_edgelist(adj_list[i]) + if max_nodes is not None and G.number_of_nodes() > max_nodes: + continue + + # add features and labels + G.graph['label'] = graph_labels[i - 1] + for u in node_iter(G): + if len(node_labels) > 0: + node_label_one_hot = [0] * num_unique_node_labels + node_label = node_labels[u - 1] + node_label_one_hot[node_label] = 1 + node_dict(G)[u]['label'] = node_label_one_hot + if len(node_attrs) > 0: + node_dict(G)[u]['feat'] = node_attrs[u - 1] + if len(node_attrs) > 0: + G.graph['feat_dim'] = node_attrs[0].shape[0] + + # relabeling + mapping = {} + it = 0 + for n in node_iter(G): + mapping[n] = it + it += 1 + # indexed from 0 + graphs.append(nx.relabel_nodes(G, mapping)) + return graphs + + +def prepare_val_data(graphs, args, val_idx, max_nodes=0): + random.shuffle(graphs) + val_size = len(graphs) // 10 + train_graphs = graphs[:val_idx * val_size] + if val_idx < 9: + train_graphs = train_graphs + graphs[(val_idx+1) * val_size :] + val_graphs = graphs[val_idx*val_size: (val_idx+1)*val_size] + + # minibatch + dataset_sampler = GraphSampler(train_graphs, normalize=False, max_num_nodes=max_nodes, + features=args.feature_type) + + train_dataset_loader = tlx.dataflow.DataLoader( + dataset_sampler, + batch_size=args.batch_size, + shuffle=True, + num_workers=args.num_workers) + + dataset_sampler = GraphSampler(val_graphs, normalize=False, max_num_nodes=max_nodes, + features=args.feature_type) + val_dataset_loader = tlx.dataflow.DataLoader( + dataset_sampler, + batch_size=args.batch_size, + shuffle=False, + num_workers=args.num_workers) + return train_dataset_loader, val_dataset_loader, \ + dataset_sampler.max_num_nodes, dataset_sampler.feat_dim, dataset_sampler.assign_feat_dim #89.89 + diff --git a/gammagl/models/diffpool.py b/gammagl/models/diffpool.py new file mode 100644 index 000000000..4dc715e86 --- /dev/null +++ b/gammagl/models/diffpool.py @@ -0,0 +1,233 @@ +import tensorlayerx as tlx +import numpy as np + +# GCN basic operation +class GraphConv(tlx.nn.Module): + @property + def all_weights(self): + return self._all_weights + + def __init__(self, input_dim, output_dim, add_self=False, normalize_embedding=False, + dropout=0.0, bias=True): + super(GraphConv, self).__init__() + self.add_self = add_self + self.dropout = dropout + if dropout > 0.001: + self.dropout_layer = tlx.nn.Dropout(p=dropout) + self.normalize_embedding = normalize_embedding + self.input_dim = input_dim + self.output_dim = output_dim + initor = tlx.initializers.TruncatedNormal() + self.weight = self._get_weights("weight", shape=(input_dim, output_dim),init=initor) + if bias: + self.bias = self._get_weights("bias", shape=(output_dim,),init=initor) + else: + self.bias = None + + def forward(self, x, adj): + if self.dropout > 0.001: + x = self.dropout_layer(x) + y = tlx.ops.matmul(adj, x) + if self.add_self: + y += x + y = tlx.ops.matmul(y, self.weight) + if self.bias is not None: + y = y + self.bias + if self.normalize_embedding: + y = tlx.ops.l2_normalize(y, axis=2) + return y + + @all_weights.setter + def all_weights(self, value): + self._all_weights = value + + +class GcnEncoderGraph(tlx.nn.Module): + r""" + Graph Convolutional Network proposed in `Semi-Supervised Classification with Graph Convolutional Networks`_. + .. _Semi-Supervised Classification with Graph Convolutional Networks: + https://arxiv.org/pdf/1609.02907.pdf + + Parameters + ---------- + input_dim (int): input dimension + hidden_dim (int): hidden dimension + embedding_dim (int): embedding dimension + label_dim (int): label dimension + num_layers (int): number of layers + pred_hidden_dims (int[]): predict hidden dimension + concat (boolean): concat to get predict input dimension + bn (boolean): batch normalization + drop_rate (float): dropout rate + args (parser): parser + """ + + @property + def all_weights(self): + return self._all_weights + + def __init__(self, input_dim, hidden_dim, embedding_dim, label_dim, num_layers, + pred_hidden_dims=[], concat=True, bn=True, dropout=0.0, args=None): + super(GcnEncoderGraph, self).__init__() + self.concat = concat + add_self = not concat + self.bn = bn + self.num_layers = num_layers + self.num_aggs = 1 + + self.bias = True + if args is not None: + self.bias = args.bias + + self.conv_first, self.conv_block, self.conv_last = self.build_conv_layers( + input_dim, hidden_dim, embedding_dim, num_layers, + add_self, normalize=True, dropout=dropout) + self.act = tlx.ReLU() + self.label_dim = label_dim + + if concat: + self.pred_input_dim = hidden_dim * (num_layers - 1) + embedding_dim + else: + self.pred_input_dim = embedding_dim + self.pred_model = self.build_pred_layers(self.pred_input_dim, pred_hidden_dims, + label_dim, num_aggs=self.num_aggs) + + if isinstance(self, GraphConv): + self.all_weights = tlx.initializers.XavierUniform(self.all_weights) + if self.bias is not None: + self.bias = tlx.initializers.constant(0.0) + + def build_conv_layers(self, input_dim, hidden_dim, embedding_dim, num_layers, add_self, + normalize=False, dropout=0.0): + conv_first = GraphConv(input_dim=input_dim, output_dim=hidden_dim, add_self=add_self, + normalize_embedding=normalize, bias=self.bias) + conv_block = tlx.nn.ModuleList( + [GraphConv(input_dim=hidden_dim, output_dim=hidden_dim, add_self=add_self, + normalize_embedding=normalize, dropout=dropout, bias=self.bias) + for i in range(num_layers - 2)]) + conv_last = GraphConv(input_dim=hidden_dim, output_dim=embedding_dim, add_self=add_self, + normalize_embedding=normalize, + bias=self.bias) + return conv_first, conv_block, conv_last + + def build_pred_layers(self, pred_input_dim, pred_hidden_dims, label_dim, num_aggs=1): + pred_input_dim = pred_input_dim * num_aggs + if len(pred_hidden_dims) == 0: + pred_model = tlx.nn.Linear(in_features=pred_input_dim, out_features=label_dim) #in&out + else: + pred_layers = [] + for pred_dim in pred_hidden_dims: + pred_layers.append(tlx.nn.Linear(in_features=pred_input_dim, out_features=pred_dim)) + pred_layers.append(self.act) + pred_input_dim = pred_dim + pred_layers.append(tlx.nn.Linear(pred_dim, label_dim)) + pred_model = tlx.nn.SequentialLayer(*pred_layers) + return pred_model + + def construct_mask(self, max_nodes, batch_num_nodes): + ''' For each num_nodes in batch_num_nodes, the first num_nodes entries of the + corresponding column are 1's, and the rest are 0's (to be masked out). + Dimension of mask: [batch_size x max_nodes x 1] + ''' + # masks + + tup= [] + for num in batch_num_nodes: + tup.append(int(num)) + packed_masks = tlx.ones(shape=(1,tup[0])) + batch_size = len(batch_num_nodes) + out_tensor = tlx.ops.zeros(shape=(batch_size, max_nodes)) + for i, mask in enumerate(packed_masks): + out_tensor[i, :batch_num_nodes[i]] = mask + return out_tensor.unsqueeze(2).cuda() + + def apply_bn(self, x): + ''' Batch normalization of 3D tensor x + ''' + device = 'cpu' + bn_module=tlx.nn.BatchNorm1d()(x.to(device)) + return bn_module + + def gcn_forward(self, x, adj, conv_first, conv_block, conv_last, embedding_mask=None): + + ''' Perform forward prop with graph convolution. + Returns: + Embedding matrix with dimension [batch_size x num_nodes x embedding] + ''' + + x = conv_first(x, adj) + x = self.act(x) + if self.bn: + x = self.apply_bn(x) + x_all = [x] + for i in range(len(conv_block)): + x = conv_block[i](x, adj) + x = self.act(x) + if self.bn: + x = self.apply_bn(x) + x_all.append(x) + x = conv_last(x, adj) + x_all.append(x) + # x_tensor: [batch_size x num_nodes x embedding] + x_tensor = tlx.ops.concat(x_all, 2) + if embedding_mask is not None: + x_tensor = x_tensor * embedding_mask + return x_tensor + + def forward(self, x, adj, batch_num_nodes=None, **kwargs): + max_num_nodes = adj.size()[1] + if batch_num_nodes is not None: + self.embedding_mask = self.construct_mask(max_num_nodes, batch_num_nodes) + else: + self.embedding_mask = None + + # conv + x = self.conv_first(x, adj) + x = self.act(x) + # print(x.is_cuda) + if self.bn: + x = self.apply_bn(x) + out_all = [] + out = tlx.ops.reduce_max(x, axis=1) + out_all.append(out) + for i in range(self.num_layers - 2): + x = self.conv_block[i](x, adj) + x = self.act(x) + if self.bn: + x = self.apply_bn(x) + out = tlx.ops.reduce_max(x, 1) + out_all.append(out) + if self.num_aggs == 2: + out = tlx.ops.cumsum(x, 1) + out_all.append(out) + x = self.conv_last(x, adj) + out= tlx.ops.reduce_max(x, 1) + out_all.append(out) + if self.num_aggs == 2: + out = tlx.ops.cumsum(x, 1) + out_all.append(out) + if self.concat: + output = tlx.ops.concat(out_all, 1) + else: + output = out + ypred = self.pred_model(output) + return ypred + + def loss(self, pred, label, type='softmax'): + # softmax + CE + if type == 'softmax': + return tlx.losses.softmax_cross_entropy_with_logits(pred, label, reduction='mean') + elif type == 'margin': + batch_size = pred.size()[0] + diag = tlx.convert_to_tensor(np.eye(tlx.reduce_max(label) + 1), dtype=tlx.int64) + label_onehot=tlx.gather(diag, label) + return tlx.losses.binary_cross_entropy(pred, label_onehot) #MultiLabelMarginLoss()(pred, label_onehot) + + @all_weights.setter + def all_weights(self, value): + self._all_weights = value + + + + +