def _one_hot_bonds(bonds: tf.Tensor) -> tf.Tensor: vocab_sizes = features.get_bond_feature_dims() one_hots = [] for i in range(bonds.shape[1]): one_hots.append( tf.one_hot(bonds[:, i], vocab_sizes[i], dtype=tf.float32)) return tf.concat(one_hots, axis=-1)
def get_tf_dataset(self): atom_feature_dims = get_atom_feature_dims() bond_feature_dims = get_bond_feature_dims() np.random.seed(23) data = { 'nodes': np.random.randint(size=(self.n_examples, self.nodes_per_graph, len(atom_feature_dims)), low=np.zeros_like(atom_feature_dims), high=atom_feature_dims, dtype=np.int32), 'edges': np.random.randint(size=(self.n_examples, self.edges_per_graph, len(bond_feature_dims)), low=np.zeros_like(bond_feature_dims), high=bond_feature_dims, dtype=np.int32), 'edge_idx': np.stack([ self.get_random_edge_idx() for _ in range(self.n_examples) ]).astype(np.int32), # binary ground truth 'ground_truth': np.random.uniform(low=0, high=2, size=(self.n_examples, 1)).astype(np.int32) } ds = tf.data.Dataset.from_tensor_slices(data) ds = ds.take(self.n_examples).cache().shuffle(32).repeat().batch( self.batch_size, drop_remainder=True) ds = ds.map(lambda b: self.batch_to_outputs(b)) ds = ds.prefetch(1024) return ds
def __init__(self, emb_dim): super().__init__() self.emb_dim = emb_dim self.bond_embedding_list = [] for dim in get_bond_feature_dims(): self.bond_embedding_list.append( tf.keras.layers.Embedding(dim, emb_dim))
def _sample_one_hot_bonds(bonds: tf.Tensor) -> tf.Tensor: vocab_sizes = features.get_bond_feature_dims() one_hots = [] num_bonds = tf.shape(bonds)[0] for i in range(bonds.shape[1]): sampled_category = _sample_uniform_categorical(num_bonds, vocab_sizes[i]) one_hots.append( tf.one_hot(sampled_category, vocab_sizes[i], dtype=tf.float32)) return tf.concat(one_hots, axis=-1)
def __init__(self, emb_dim): super(ExampleEdgeEncoder, self).__init__() self.bond_embedding_list = torch.nn.ModuleList() full_bond_feature_dims = get_bond_feature_dims() for i, dim in enumerate(full_bond_feature_dims): emb = torch.nn.Embedding(dim, emb_dim) torch.nn.init.xavier_uniform_(emb.weight.data) self.bond_embedding_list.append(emb)
def __init__(self, emb_dim): super().__init__() from ogb.utils.features import get_bond_feature_dims self.bond_embedding_list = torch.nn.ModuleList() for i, dim in enumerate(get_bond_feature_dims()): emb = torch.nn.Embedding(dim, emb_dim) torch.nn.init.xavier_uniform_(emb.weight.data) self.bond_embedding_list.append(emb)
def __call__(self, x): bond_feature = get_bond_feature_dims() bond_input = L.split(x, num_or_sections=len(bond_feature), dim=-1) outputs = None count = 0 for _x, _bond_input_dim in zip(bond_input, bond_feature): count += 1 emb = L.embedding(_x, size=(_bond_input_dim, self.emb_dim), param_attr=F.ParamAttr(name=self.name + '_bond_feat_%s' % count)) if outputs is None: outputs = emb else: outputs = outputs + emb return outputs
import torch from ogb.utils.features import get_atom_feature_dims, get_bond_feature_dims full_atom_feature_dims = get_atom_feature_dims() full_bond_feature_dims = get_bond_feature_dims() class AtomEncoder(torch.nn.Module): def __init__(self, emb_dim): super(AtomEncoder, self).__init__() self.atom_embedding_list = torch.nn.ModuleList() for i, dim in enumerate(full_atom_feature_dims): emb = torch.nn.Embedding(dim, emb_dim) torch.nn.init.xavier_uniform_(emb.weight.data) self.atom_embedding_list.append(emb) def forward(self, x): x_embedding = 0 for i in range(x.shape[1]): x_embedding += self.atom_embedding_list[i](x[:, i]) return x_embedding class BondEncoder(torch.nn.Module): def __init__(self, emb_dim): super(BondEncoder, self).__init__() self.bond_embedding_list = torch.nn.ModuleList()
import torch import torch.nn as nn from typing import Union from phc.quaternion.algebra import QTensor from phc.quaternion.activations import get_functional_activation from phc.quaternion.layers import quaternion_dropout, QLinear, RealTransformer from phc.quaternion.norm import QuaternionNorm from ogb.utils.features import get_atom_feature_dims, get_bond_feature_dims ATOM_FEAT_DIMS = get_atom_feature_dims() BOND_FEAT_DIMS = get_bond_feature_dims() """ Quaternion Downstream Feed-Forward Network""" class QuaternionDownstreamNet(nn.Module): """ A quaternion Feed-Forward Network which predicts a real-valued vector of dimension `out_features`. """ def __init__(self, in_features: int, hidden_layers: list, out_features: int, activation: str, bias: bool, norm: str, init: str, dropout: Union[float, list], same_dropout: bool = False, real_trafo: str = "linear") -> None: super(QuaternionDownstreamNet, self).__init__()
# Copyright (c) 2021 Graphcore Ltd. All rights reserved. import numpy as np import tensorflow as tf from absl import flags from ogb.utils.features import get_atom_feature_dims, get_bond_feature_dims NODE_FEATURE_DIMS = len(get_atom_feature_dims()) EDGE_FEATURE_DIMS = len(get_bond_feature_dims()) FLAGS = flags.FLAGS def np_batch_generator(n_nodes, n_edges, batch_size, data_subset, epochs=1, shuffle=True): """ generates the batches in numpy shape of batch is going to be: nodes: [max_nodes, n_node_feat] edge_idx: [max_edges, 2] edge_feat: [max_edges, n_edge_feat] labels: [max_graphs, 1] :param n_nodes: the number of nodes per graph (on average) :param n_edges: the number of edges per graph (on average) :param batch_size: desired batch size (1 is reserved for the dummy graph)
import unittest import torch import itertools from ogb.utils.features import get_atom_feature_dims, get_bond_feature_dims ATOM_FEATS, BOND_FEATS = get_atom_feature_dims(), get_bond_feature_dims() from phc.quaternion.layers import QLinear from phc.quaternion.encoder import IntegerEncoder, QuaternionEncoder # quaternion undirectional models from phc.quaternion.undirectional.models import QuaternionSkipConnectAdd as UQ_SC_ADD from phc.quaternion.undirectional.models import QuaternionSkipConnectConcat as UQ_SC_CAT def weights_flag(name): """ Select only the weight matrices of QLinear and nn.Embedding """ return "W_" in name or "weight" in name and "b_" not in name and "bias" not in name and "bn" not in name def check_qlinear(in_features: int, out_features: int, init: str, norm_tol: float = 5.0, nruns: int = 10) -> bool: weights = [] module = QLinear(in_features, out_features, bias=False, init=init) for i in range(nruns): weights.append(module.W.stack(dim=0)) module.reset_parameters()
def main(): args = get_parser() # get some argparse arguments that are parsed a bool string naive_encoder = not str2bool(args.full_encoder) pin_memory = str2bool(args.pin_memory) use_bias = str2bool(args.bias) downstream_bn = str(args.d_bn) same_dropout = str2bool(args.same_dropout) mlp_mp = str2bool(args.mlp_mp) phm_dim = args.phm_dim learn_phm = str2bool(args.learn_phm) base_dir = "pcba/" if not os.path.exists(base_dir): os.makedirs(base_dir) if base_dir not in args.save_dir: args.save_dir = os.path.join(base_dir, args.save_dir) if not os.path.exists(args.save_dir): os.makedirs(args.save_dir) set_logging(save_dir=args.save_dir) logging.info(f"Creating log directory at {args.save_dir}.") with open(os.path.join(args.save_dir, "params.json"), 'w') as fp: json.dump(args.__dict__, fp) mp_layers = [int(item) for item in args.mp_units.split(',')] downstream_layers = [int(item) for item in args.d_units.split(',')] mp_dropout = [float(item) for item in args.dropout_mpnn.split(',')] dn_dropout = [float(item) for item in args.dropout_dn.split(',')] logging.info( f'Initialising model with {mp_layers} hidden units with dropout {mp_dropout} ' f'and downstream units: {downstream_layers} with dropout {dn_dropout}.' ) if args.pooling == "globalsum": logging.info("Using GlobalSum Pooling") else: logging.info("Using SoftAttention Pooling") logging.info( f"Using Adam optimizer with weight_decay ({args.weightdecay}) and regularization " f"norm ({args.regularization})") logging.info( f"Weight init: {args.w_init} \n Contribution init: {args.c_init}") # data dname = "ogbg-molpcba" transform = RemoveIsolatedNodes() # pre-transform doesnt work somehow.. dataset = PygGraphPropPredDataset( name=dname, root="dataset") #, pre_transform=transform, transform=None) evaluator = Evaluator(name=dname) split_idx = dataset.get_idx_split() train_data = dataset[split_idx["train"]] valid_data = dataset[split_idx["valid"]] test_data = dataset[split_idx["test"]] if PRE_TRAFO: # pre-transform in memory to overcome computations when training logging.info( "Pre-transforming graphs, to overcome computation in batching.") train_data_list = [] valid_data_list = [] test_data_list = [] for data in train_data: train_data_list.append(transform(data)) for data in valid_data: valid_data_list.append(transform(data)) for data in test_data: test_data_list.append(transform(data)) logging.info("finised. Initiliasing dataloaders") train_loader = DataLoader(train_data_list, batch_size=args.batch_size, drop_last=False, shuffle=True, num_workers=args.nworkers, pin_memory=pin_memory) valid_loader = DataLoader(valid_data_list, batch_size=args.batch_size, drop_last=False, shuffle=False, num_workers=args.nworkers, pin_memory=pin_memory) test_loader = DataLoader(test_data_list, batch_size=args.batch_size, drop_last=False, shuffle=False, num_workers=args.nworkers, pin_memory=pin_memory) else: train_loader = DataLoader(train_data, batch_size=args.batch_size, drop_last=False, shuffle=True, num_workers=args.nworkers, pin_memory=pin_memory) valid_loader = DataLoader(valid_data, batch_size=args.batch_size, drop_last=False, shuffle=False, num_workers=args.nworkers, pin_memory=pin_memory) test_loader = DataLoader(test_data, batch_size=args.batch_size, drop_last=False, shuffle=False, num_workers=args.nworkers, pin_memory=pin_memory) device = f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu' FULL_ATOM_FEATURE_DIMS = get_atom_feature_dims() FULL_BOND_FEATURE_DIMS = get_bond_feature_dims() # for hypercomplex model unique_phm = str2bool(args.unique_phm) if unique_phm: phm_rule = get_multiplication_matrices(phm_dim=args.phm_dim, type="phm") phm_rule = torch.nn.ParameterList( [torch.nn.Parameter(a, requires_grad=learn_phm) for a in phm_rule]) else: phm_rule = None if args.aggr_msg == "pna" or args.aggr_node == "pna": # if PNA is used # Compute in-degree histogram over training data. deg = torch.zeros(6, dtype=torch.long) for data in dataset[split_idx['train']]: d = degree(data.edge_index[1], num_nodes=data.num_nodes, dtype=torch.long) deg += torch.bincount(d, minlength=deg.numel()) else: deg = None aggr_kwargs = { "aggregators": ['mean', 'min', 'max', 'std'], "scalers": ['identity', 'amplification', 'attenuation'], "deg": deg, "post_layers": 1, "msg_scalers": str2bool(args.msg_scale ), # this key is for directional messagepassing layers. "initial_beta": 1.0, # Softmax "learn_beta": True } if "quaternion" in args.type: if args.aggr_msg == "pna" or args.aggr_msg == "pna": logging.info("PNA not implemented for quaternion models.") raise NotImplementedError if args.type == "undirectional-quaternion-sc-add": logging.info( "Using Quaternion Undirectional MPNN with Skip Connection through Addition" ) model = UQ_SC_ADD(atom_input_dims=FULL_ATOM_FEATURE_DIMS, atom_encoded_dim=args.input_embed_dim, bond_input_dims=FULL_BOND_FEATURE_DIMS, naive_encoder=naive_encoder, mp_layers=mp_layers, dropout_mpnn=mp_dropout, init=args.w_init, same_dropout=same_dropout, norm_mp=args.mp_norm, add_self_loops=True, msg_aggr=args.aggr_msg, node_aggr=args.aggr_node, mlp=mlp_mp, pooling=args.pooling, activation=args.activation, real_trafo=args.real_trafo, downstream_layers=downstream_layers, target_dim=dataset.num_tasks, dropout_dn=dn_dropout, norm_dn=downstream_bn, msg_encoder=args.msg_encoder, **aggr_kwargs) elif args.type == "undirectional-quaternion-sc-cat": logging.info( "Using Quaternion Undirectional MPNN with Skip Connection through Concatenation" ) model = UQ_SC_CAT(atom_input_dims=FULL_ATOM_FEATURE_DIMS, atom_encoded_dim=args.input_embed_dim, bond_input_dims=FULL_BOND_FEATURE_DIMS, naive_encoder=naive_encoder, mp_layers=mp_layers, dropout_mpnn=mp_dropout, init=args.w_init, same_dropout=same_dropout, norm_mp=args.mp_norm, add_self_loops=True, msg_aggr=args.aggr_msg, node_aggr=args.aggr_node, mlp=mlp_mp, pooling=args.pooling, activation=args.activation, real_trafo=args.real_trafo, downstream_layers=downstream_layers, target_dim=dataset.num_tasks, dropout_dn=dn_dropout, norm_dn=downstream_bn, msg_encoder=args.msg_encoder, **aggr_kwargs) elif args.type == "undirectional-phm-sc-add": logging.info( "Using PHM Undirectional MPNN with Skip Connection through Addition" ) model = UPH_SC_ADD(phm_dim=phm_dim, learn_phm=learn_phm, phm_rule=phm_rule, atom_input_dims=FULL_ATOM_FEATURE_DIMS, atom_encoded_dim=args.input_embed_dim, bond_input_dims=FULL_BOND_FEATURE_DIMS, naive_encoder=naive_encoder, mp_layers=mp_layers, dropout_mpnn=mp_dropout, w_init=args.w_init, c_init=args.c_init, same_dropout=same_dropout, norm_mp=args.mp_norm, add_self_loops=True, msg_aggr=args.aggr_msg, node_aggr=args.aggr_node, mlp=mlp_mp, pooling=args.pooling, activation=args.activation, real_trafo=args.real_trafo, downstream_layers=downstream_layers, target_dim=dataset.num_tasks, dropout_dn=dn_dropout, norm_dn=downstream_bn, msg_encoder=args.msg_encoder, sc_type=args.sc_type, **aggr_kwargs) elif args.type == "undirectional-phm-sc-cat": logging.info( "Using PHM Undirectional MPNN with Skip Connection through Concatenation" ) model = UPH_SC_CAT(phm_dim=phm_dim, learn_phm=learn_phm, phm_rule=phm_rule, atom_input_dims=FULL_ATOM_FEATURE_DIMS, atom_encoded_dim=args.input_embed_dim, bond_input_dims=FULL_BOND_FEATURE_DIMS, naive_encoder=naive_encoder, mp_layers=mp_layers, dropout_mpnn=mp_dropout, w_init=args.w_init, c_init=args.c_init, same_dropout=same_dropout, norm_mp=args.mp_norm, add_self_loops=True, msg_aggr=args.aggr_msg, node_aggr=args.aggr_node, mlp=mlp_mp, pooling=args.pooling, activation=args.activation, real_trafo=args.real_trafo, downstream_layers=downstream_layers, target_dim=dataset.num_tasks, dropout_dn=dn_dropout, norm_dn=downstream_bn, msg_encoder=args.msg_encoder, **aggr_kwargs) else: raise ModuleNotFoundError logging.info( f"Model consists of {model.get_number_of_params_()} trainable parameters" ) # do runs test_best_epoch_metrics_arr = [] test_last_epoch_metrics_arr = [] val_metrics_arr = [] for i in range(1, args.n_runs + 1): ogb_bestEpoch_test_metrics, ogb_lastEpoch_test_metric, ogb_val_metrics = do_run( i, model, args, transform, train_loader, valid_loader, test_loader, device, evaluator) test_best_epoch_metrics_arr.append(ogb_bestEpoch_test_metrics) test_last_epoch_metrics_arr.append(ogb_lastEpoch_test_metric) val_metrics_arr.append(ogb_val_metrics) logging.info(f"Performance of model across {args.n_runs} runs:") test_bestEpoch_perf = torch.tensor(test_best_epoch_metrics_arr) test_lastEpoch_perf = torch.tensor(test_last_epoch_metrics_arr) valid_perf = torch.tensor(val_metrics_arr) logging.info('===========================') logging.info( f'Final Test (best val-epoch) ' f'"{evaluator.eval_metric}": {test_bestEpoch_perf.mean():.4f} ± {test_bestEpoch_perf.std():.4f}' ) logging.info( f'Final Test (last-epoch) ' f'"{evaluator.eval_metric}": {test_lastEpoch_perf.mean():.4f} ± {test_lastEpoch_perf.std():.4f}' ) logging.info( f'Final (best) Valid "{evaluator.eval_metric}": {valid_perf.mean():.4f} ± {valid_perf.std():.4f}' )
def __init__(self, encoder_name, d_in_features, d_in_encoder, d_out_encoder, **kwargs): super(DiscreteEmbedding, self).__init__() #-------------- various different embedding layers kwargs['init'] = None if 'init' not in kwargs else kwargs['init'] self.encoder_name = encoder_name # d_in_features: input feature size (e.g. if already one hot encoded), # d_in_encoder: number of unique values that will be encoded (size of embedding vocabulary) #-------------- fill embedding with zeros if encoder_name == 'zero_encoder': self.encoder = zero_encoder(d_out_encoder) d_out = d_out_encoder #-------------- linear pojection elif encoder_name == 'linear': self.encoder = nn.Linear(d_in_features, d_out_encoder, bias=True) d_out = d_out_encoder #-------------- mlp elif encoder_name == 'mlp': self.encoder = mlp(d_in_features, d_out_encoder, d_out_encoder, kwargs['seed'], kwargs['activation_mlp'], kwargs['bn_mlp']) d_out = d_out_encoder #-------------- multi hot encoding of categorical data elif encoder_name == 'one_hot_encoder': self.encoder = one_hot_encoder(d_in_encoder) d_out = sum(d_in_encoder) #-------------- embedding of categorical data (linear projection without bias of one hot encodings) elif encoder_name == 'embedding': self.encoder = multi_embedding(d_in_encoder, d_out_encoder, kwargs['aggr'], kwargs['init']) if kwargs['aggr'] == 'concat': d_out = len(d_in_encoder) * d_out_encoder else: d_out = d_out_encoder #-------------- for ogb: multi hot encoding of node features elif encoder_name == 'atom_one_hot_encoder': full_atom_feature_dims = get_atom_feature_dims( ) if kwargs['features_scope'] == 'full' else get_atom_feature_dims( )[:2] self.encoder = one_hot_encoder(full_atom_feature_dims) d_out = sum(full_atom_feature_dims) #-------------- for ogb: multi hot encoding of edge features elif encoder_name == 'bond_one_hot_encoder': full_bond_feature_dims = get_bond_feature_dims( ) if kwargs['features_scope'] == 'full' else get_bond_feature_dims( )[:2] self.encoder = one_hot_encoder(full_bond_feature_dims) d_out = sum(full_bond_feature_dims) #-------------- for ogb: embedding of node features elif encoder_name == 'atom_encoder': self.encoder = AtomEncoder(d_out_encoder) d_out = d_out_encoder #-------------- for ogb: embedding of edge features elif encoder_name == 'bond_encoder': self.encoder = BondEncoder(emb_dim=d_out_encoder) d_out = d_out_encoder #-------------- no embedding, use as is elif encoder_name == 'None': self.encoder = None d_out = d_in_features else: raise NotImplementedError( 'Encoder {} is not currently supported.'.format(encoder_name)) self.d_out = d_out return