[Model] EGT (#238)

* EGT模型

* add egt

* add egt

---------

Co-authored-by: Tan Yingying <[email protected]>
Co-authored-by: Xingyuan Ji <[email protected]>

Author: 3 people

examples/egt/__init__.py

@@ -0,0 +1,23 @@
+scheme: pcqm4mv2
+model_name: egt_90m
+distributed: false         # 设置为true启用多GPU分布式训练
+batch_size: 512             # 单GPU时的批次大小;多GPU时需调整为总批次大小除以GPU数量
+model_height: 24
+node_width: 768
+edge_width: 64
+num_heads: 32
+num_epochs: 1
+max_lr: 0.0001
+attn_dropout: 0.3
+lr_warmup_steps: 200000
+lr_total_steps: 1000000
+node_ffn_multiplier: 1.0
+edge_ffn_multiplier: 1.0
+upto_hop: 16
+dataloader_workers: 1       # 多线程数
+scale_degree: true
+num_virtual_nodes: 4
+svd_random_neg: true
+mixed_precision: true       # 启用混合精度
+use_adaptive_sparse: true   # 启用自适应稀疏
+sparse_alpha: 0.5           # 稀疏化强度系数

examples/egt/config.yaml

@@ -0,0 +1,106 @@
+import numpy as np
+import tensorlayerx as tlx
+from tqdm import tqdm
+from pathlib import Path
+
+
+class DatasetBase:
+    def __init__(self,
+                 dataset_name,
+                 split,
+                 cache_dir=None,
+                 load_cache_if_exists=True,
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.dataset_name = dataset_name
+        self.split = split
+        self.cache_dir = cache_dir
+
+        self.is_cached = False
+        if load_cache_if_exists:
+            self.cache(verbose=0, must_exist=True)
+
+    @property
+    def record_tokens(self):
+        raise NotImplementedError
+
+    def read_record(self, token):
+        raise NotImplementedError
+
+    def __len__(self):
+        return len(self.record_tokens)
+
+    def __getitem__(self, index):
+
+        token = self.record_tokens[index]
+        try:
+            return self._records[token]
+        except AttributeError:
+            record = self.read_record(token)
+            self._records = {token: record}
+            return record
+        except KeyError:
+            record = self.read_record(token)
+            self._records[token] = record
+            return record
+
+    def read_all_records(self, verbose=1):
+        self._records = {}
+        if verbose:
+            print(f'Reading all {self.split} records...', flush=True)
+            for token in tqdm(self.record_tokens):
+                self._records[token] = self.read_record(token)
+        else:
+            for token in self.record_tokens:
+                self._records[token] = self.read_record(token)
+
+    def get_cache_path(self, path=None):
+        if path is None:
+            path = self.cache_dir
+        base_path = (Path(path)/self.dataset_name)/self.split
+        base_path.mkdir(parents=True, exist_ok=True)
+        return base_path
+
+    def cache_load_and_save(self, base_path, op, verbose):
+        tokens_path = base_path/'tokens.npy'
+        records_path = base_path/'records.npy'
+
+        if op == 'load':
+            self._record_tokens = tlx.files.load_npy_to_any(
+                name=str(tokens_path))
+            self._records = tlx.files.load_npy_to_any(name=str(records_path))
+        elif op == 'save':
+            if tokens_path.exists() and records_path.exists() \
+                    and hasattr(self, '_record_tokens') and hasattr(self, '_records'):
+                return
+            self.read_all_records(verbose=verbose)
+            tlx.files.save_any_to_npy(
+                save_dict=self.record_tokens, name=str(tokens_path))
+            tlx.files.save_any_to_npy(
+                save_dict=self._records, name=str(records_path))
+        else:
+            raise ValueError(f'Unknown operation: {op}')
+
+    def cache(self, path=None, verbose=1, must_exist=False):
+        if self.is_cached:
+            return
+
+        base_path = self.get_cache_path(path)
+        try:
+            if verbose:
+                print(
+                    f'Trying to load {self.split} cache from disk...', flush=True)
+            self.cache_load_and_save(base_path, 'load', verbose)
+            if verbose:
+                print(f'Loaded {self.split} cache from disk.', flush=True)
+        except FileNotFoundError:
+            if must_exist:
+                return
+
+            if verbose:
+                print(f'{self.split} cache does not exist! Cacheing...', flush=True)
+            self.cache_load_and_save(base_path, 'save', verbose)
+            if verbose:
+                print(f'Saved {self.split} cache to disk.', flush=True)
+
+        self.is_cached = True

examples/egt/data/dataset_base.py

@@ -0,0 +1,97 @@
+import tensorlayerx as tlx  
+import numpy as np
+
+from .dataset_base import DatasetBase
+
+from .stack_with_pad import stack_with_pad
+from collections import defaultdict
+from numba.typed import List
+
+
+class GraphDataset(DatasetBase):
+    def __init__(self,
+                 num_nodes_key='num_nodes',
+                 edges_key='edges',
+                 node_features_key='node_features',
+                 edge_features_key='edge_features',
+                 node_mask_key='node_mask',
+                 targets_key='target',
+                 include_node_mask=True,
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.num_nodes_key = num_nodes_key
+        self.edges_key = edges_key
+        self.node_features_key = node_features_key
+        self.edge_features_key = edge_features_key
+        self.node_mask_key = node_mask_key
+        self.targets_key = targets_key
+        self.include_node_mask = include_node_mask
+
+
+
+    def __getitem__(self, index):
+        item = super().__getitem__(index)
+        if self.include_node_mask:
+            item = item.copy()
+            item[self.node_mask_key] = np.ones((item[self.num_nodes_key],), dtype=np.uint8)
+        return item
+
+    def _calculate_max_nodes(self):
+        self._max_nodes = self[0][self.num_nodes_key]
+        self._max_nodes_index = 0
+        for i in range(1, super().__len__()):
+            graph = super().__getitem__(i)
+            cur_nodes = graph[self.num_nodes_key]
+            if cur_nodes > self._max_nodes:
+                self._max_nodes = cur_nodes
+                self._max_nodes_index = i
+
+    @property
+    def max_nodes(self):
+        try:
+            return self._max_nodes
+        except AttributeError:
+            self._calculate_max_nodes()
+            return self._max_nodes
+
+    @property
+    def max_nodes_index(self):
+        try:
+            return self._max_nodes_index
+        except AttributeError:
+            self._calculate_max_nodes()
+            return self._max_nodes_index
+
+    def cache_load_and_save(self, base_path, op, verbose):
+        super().cache_load_and_save(base_path, op, verbose)
+        max_nodes_path = base_path / 'max_nodes_data.npy' 
+
+        if op == 'load':
+            # 替换 torch.load 为 np.load
+            max_nodes_data = tlx.files.load_npy_to_any(name=str(max_nodes_path))
+            self._max_nodes = max_nodes_data['max_nodes']
+            self._max_nodes_index = max_nodes_data['max_nodes_index']
+        elif op == 'save':
+            if verbose:
+                print(f'Calculating {self.split} max nodes...', flush=True)
+            max_nodes_data = {'max_nodes': self.max_nodes,
+                              'max_nodes_index': self.max_nodes_index}
+            # 替换 torch.save 为 np.savez
+            tlx.files.save_any_to_npy(save_dict=max_nodes_data, name=str(max_nodes_path))
+        else:
+            raise ValueError(f'Unknown operation: {op}')
+
+    def max_batch(self, batch_size, collate_fn):
+        return collate_fn([self.__getitem__(self.max_nodes_index)] * batch_size)
+
+
+def graphdata_collate(batch):
+    batch_data = defaultdict(List)
+    for elem in batch:
+        for k, v in elem.items():
+            batch_data[k].append(v)
+
+    # 替换 torch.from_numpy 为 tlx.convert_to_tensor
+    out = {k: tlx.convert_to_tensor(stack_with_pad(dat))
+           for k, dat in batch_data.items()}
+    return out

examples/egt/data/graph_dataset.py

@@ -0,0 +1,91 @@
+import numpy as np
+import numba as nb
+
+
+@nb.njit
+def stack_with_pad_4d(inputs):
+    num_elem = len(inputs)
+    ms_0, ms_1, ms_2, ms_3 = inputs[0].shape
+    
+    for i in range(1,num_elem):
+        is_0, is_1, is_2, is_3 = inputs[i].shape
+        ms_0 = max(is_0, ms_0)
+        ms_1 = max(is_1, ms_1)
+        ms_2 = max(is_2, ms_2)
+        ms_3 = max(is_3, ms_3)
+    
+    stacked_shape = (num_elem,ms_0,ms_1,ms_2,ms_3)
+    stacked = np.zeros(stacked_shape, dtype=inputs[0].dtype)
+    
+    for i, elem in enumerate(inputs):
+        stacked[i][:elem.shape[0],:elem.shape[1],:elem.shape[2],:elem.shape[3]] = elem
+    return stacked
+
+@nb.njit
+def stack_with_pad_3d(inputs):
+    num_elem = len(inputs)
+    ms_0, ms_1, ms_2 = inputs[0].shape
+    
+    for i in range(1,num_elem):
+        is_0, is_1, is_2 = inputs[i].shape
+        ms_0 = max(is_0, ms_0)
+        ms_1 = max(is_1, ms_1)
+        ms_2 = max(is_2, ms_2)
+    
+    stacked_shape = (num_elem,ms_0,ms_1,ms_2)
+    stacked = np.zeros(stacked_shape, dtype=inputs[0].dtype)
+    
+    for i, elem in enumerate(inputs):
+        stacked[i][:elem.shape[0],:elem.shape[1],:elem.shape[2]] = elem
+    return stacked
+
+@nb.njit
+def stack_with_pad_2d(inputs):
+    num_elem = len(inputs)
+    ms_0, ms_1 = inputs[0].shape
+    
+    for i in range(1,num_elem):
+        is_0, is_1 = inputs[i].shape
+        ms_0 = max(is_0, ms_0)
+        ms_1 = max(is_1, ms_1)
+    
+    stacked_shape = (num_elem,ms_0,ms_1)
+    stacked = np.zeros(stacked_shape, dtype=inputs[0].dtype)
+    
+    for i, elem in enumerate(inputs):
+        stacked[i][:elem.shape[0],:elem.shape[1]] = elem
+    return stacked
+
+@nb.njit
+def stack_with_pad_1d(inputs):
+    num_elem = len(inputs)
+    ms_0 = inputs[0].shape[0]
+    
+    for i in range(1,num_elem):
+        is_0 = inputs[i].shape[0]
+        ms_0 = max(is_0, ms_0)
+    
+    stacked_shape = (num_elem,ms_0)
+    stacked = np.zeros(stacked_shape, dtype=inputs[0].dtype)
+    
+    for i, elem in enumerate(inputs):
+        stacked[i][:elem.shape[0]] = elem
+    return stacked
+
+
+def stack_with_pad(inputs):
+    shape_rank = np.ndim(inputs[0])
+    if shape_rank == 0:
+        return np.stack(inputs)
+    if shape_rank == 1:
+        return stack_with_pad_1d(inputs)
+    elif shape_rank == 2:
+        return stack_with_pad_2d(inputs)
+    elif shape_rank == 3:
+        return stack_with_pad_3d(inputs)
+    elif shape_rank == 4:
+        return stack_with_pad_4d(inputs)
+    else:
+        raise ValueError('Only support up to 4D tensor')
+
+

examples/egt/data/stack_with_pad.py

@@ -0,0 +1,73 @@
+import numpy as np
+import numba as nb
+
+from .graph_dataset import GraphDataset
+
+NODE_FEATURES_OFFSET = 128
+EDGE_FEATURES_OFFSET = 8
+
+@nb.njit
+def floyd_warshall(A):
+    n = A.shape[0]
+    D = np.zeros((n,n), dtype=np.int16)
+    
+    for i in range(n):
+        for j in range(n):
+            if i == j:
+                pass
+            elif A[i,j] == 0:
+                D[i,j] = 510
+            else:
+                D[i,j] = 1
+    
+    for k in range(n):
+        for i in range(n):
+            for j in range(n):
+                old_dist = D[i,j]
+                new_dist = D[i,k] + D[k,j]
+                if new_dist < old_dist:
+                    D[i,j] = new_dist
+    return D
+
+@nb.njit
+def preprocess_data(num_nodes, edges, node_feats, edge_feats):
+    node_feats = node_feats + np.arange(1,node_feats.shape[-1]*NODE_FEATURES_OFFSET+1,
+                                            NODE_FEATURES_OFFSET,dtype=np.int16)
+    edge_feats = edge_feats + np.arange(1,edge_feats.shape[-1]*EDGE_FEATURES_OFFSET+1,
+                                            EDGE_FEATURES_OFFSET,dtype=np.int16)
+    
+    A = np.zeros((num_nodes,num_nodes),dtype=np.int16)
+    E = np.zeros((num_nodes,num_nodes,edge_feats.shape[-1]),dtype=np.int16)
+    for k in range(edges.shape[0]):
+        i,j = edges[k,0], edges[k,1]
+        A[i,j] = 1
+        E[i,j] = edge_feats[k]
+    
+    D = floyd_warshall(A)
+    return node_feats, D, E
+
+
+class StructuralDataset(GraphDataset):
+    def __init__(self,
+                 distance_matrix_key      = 'distance_matrix',
+                 feature_matrix_key       = 'feature_matrix',
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.distance_matrix_key      = distance_matrix_key      
+        self.feature_matrix_key       = feature_matrix_key  
+        
+    def __getitem__(self, index):
+        item = super().__getitem__(index)
+        
+        num_nodes = int(item[self.num_nodes_key])
+        edges = item.pop(self.edges_key)
+        node_feats = item.pop(self.node_features_key)
+        edge_feats = item.pop(self.edge_features_key)
+        
+        node_feats, dist_mat, edge_feats_mat = preprocess_data(num_nodes, edges, node_feats, edge_feats)
+        item[self.node_features_key] = node_feats
+        item[self.distance_matrix_key] = dist_mat
+        item[self.feature_matrix_key] = edge_feats_mat
+        
+        return item
+