def make_inorp( # Input input_node_shape, input_edge_shape, input_state_shape, input_embedd: dict = None, # Output output_embedd: dict = None, output_mlp: dict = None, # Model specific parameter depth=3, use_set2set: bool = False, # not in original paper node_mlp_args: dict = None, edge_mlp_args: dict = None, set2set_args: dict = None, pooling_args: dict = None): """ Generate Interaction network. Args: input_node_shape (list): Shape of node features. If shape is (None,) embedding layer is used. input_edge_shape (list): Shape of edge features. If shape is (None,) embedding layer is used. input_state_shape (list): Shape of state features. If shape is (,) embedding layer is used. input_embedd (dict): Dictionary of embedding parameters used if input shape is None. Default is {'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100, 'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64, 'input_type': 'ragged'}. output_embedd (dict): Dictionary of embedding parameters of the graph network. Default is {"output_mode": 'graph', "output_type": 'padded'}. output_mlp (dict): Dictionary of arguments for final MLP regression or classifcation layer. Default is {"use_bias": [True, True, False], "units": [25, 10, 1], "activation": ['relu', 'relu', 'sigmoid']}. depth (int): Number of convolution layers. Default is 3. node_mlp_args (dict): Dictionary of arguments for MLP for node update. Default is {"units": [100, 50], "use_bias": True, "activation": ['relu', "linear"]} edge_mlp_args (dict): Dictionary of arguments for MLP for interaction update. Default is {"units": [100, 100, 100, 100, 50], "activation": ['relu', 'relu', 'relu', 'relu', "linear"]} use_set2set (str): Use set2set pooling for graph embedding. Default is False. set2set_args (dict): Dictionary of set2set layer arguments. Default is {'channels': 32, 'T': 3, "pooling_method": "mean", "init_qstar": "mean"}. pooling_args (dict): Dictionary for message pooling arguments. Default is {'is_sorted': False, 'has_unconnected': True, 'pooling_method': "segment_mean"} Returns: model (tf.keras.model): Interaction model. """ # default values model_default = { 'input_embedd': { 'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100, 'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64, 'input_tensor_type': 'ragged' }, 'output_embedd': { "output_mode": 'graph', "output_tensor_type": 'padded' }, 'output_mlp': { "use_bias": [True, True, False], "units": [25, 10, 1], "activation": ['relu', 'relu', 'sigmoid'] }, 'set2set_args': { 'channels': 32, 'T': 3, "pooling_method": "mean", "init_qstar": "mean" }, 'node_mlp_args': { "units": [100, 50], "use_bias": True, "activation": ['relu', "linear"] }, 'edge_mlp_args': { "units": [100, 100, 100, 100, 50], "activation": ['relu', 'relu', 'relu', 'relu', "linear"] }, 'pooling_args': { 'is_sorted': False, 'has_unconnected': True, 'pooling_method': "segment_mean" } } # Update default values input_embedd = update_model_args(model_default['input_embedd'], input_embedd) output_embedd = update_model_args(model_default['output_embedd'], output_embedd) output_mlp = update_model_args(model_default['output_mlp'], output_mlp) set2set_args = update_model_args(model_default['set2set_args'], set2set_args) node_mlp_args = update_model_args(model_default['node_mlp_args'], node_mlp_args) edge_mlp_args = update_model_args(model_default['edge_mlp_args'], edge_mlp_args) pooling_args = update_model_args(model_default['pooling_args'], pooling_args) gather_args = {"node_indexing": "sample"} # Make input embedding, if no feature dimension node_input, n, edge_input, ed, edge_index_input, env_input, uenv = generate_standard_graph_input( input_node_shape, input_edge_shape, input_state_shape, **input_embedd) # Preprocessing edi = edge_index_input ev = GatherState(**gather_args)([uenv, n]) # n-Layer Step for i in range(0, depth): # upd = GatherNodes()([n,edi]) eu1 = GatherNodesIngoing(**gather_args)([n, edi]) eu2 = GatherNodesOutgoing(**gather_args)([n, edi]) upd = Concatenate(axis=-1)([eu2, eu1]) eu = Concatenate(axis=-1)([upd, ed]) eu = MLP(**edge_mlp_args)(eu) # Pool message nu = PoolingLocalEdges(**pooling_args)( [n, eu, edi]) # Summing for each node connection # Add environment nu = Concatenate(axis=-1)( [n, nu, ev]) # Concatenate node features with new edge updates n = MLP(**node_mlp_args)(nu) if output_embedd["output_mode"] == 'graph': if use_set2set: # output outss = Dense(set2set_args["channels"], activation="linear")(n) out = Set2Set(**set2set_args)(outss) else: out = PoolingNodes(**pooling_args)(n) output_mlp.update({"input_tensor_type": "tensor"}) main_output = MLP(**output_mlp)(out) else: # Node labeling out = n main_output = MLP(**output_mlp)(out) main_output = ChangeTensorType( input_tensor_type="ragged", output_tensor_type="tensor")(main_output) # no ragged for distribution atm model = ks.models.Model( inputs=[node_input, edge_input, edge_index_input, env_input], outputs=main_output) return model
def make_gat( # Input input_node_shape, input_edge_shape, input_embedd: dict = None, # Output output_embedd: dict = None, output_mlp: dict = None, # Model specific parameter depth=3, attention_heads_num=5, attention_heads_concat=False, attention_args: dict = None): """ Generate Interaction network. Args: input_node_shape (list): Shape of node features. If shape is (None,) embedding layer is used. input_edge_shape (list): Shape of edge features. If shape is (None,) embedding layer is used. input_embedd (dict): Dictionary of embedding parameters used if input shape is None. Default is {'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100, 'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64, 'input_type': 'ragged'}. output_embedd (dict): Dictionary of embedding parameters of the graph network. Default is {"output_mode": 'graph', "output_type": 'padded'}. output_mlp (dict): Dictionary of arguments for final MLP regression or classifcation layer. Default is {"use_bias": [True, True, False], "units": [25, 10, 1], "activation": ['relu', 'relu', 'sigmoid']}. depth (int): Number of convolution layers. Default is 3. attention_heads_num (int): Number of attention heads. Default is 5. attention_heads_concat (bool): Concat attention. Default is False. attention_args (dict): Layer arguments for attention layer. Default is {"units": 32, 'is_sorted': False, 'has_unconnected': True} Returns: model (tf.keras.model): Interaction model. """ # default values model_default = { 'input_embedd': { 'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100, 'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64, 'input_tensor_type': 'ragged' }, 'output_embedd': { "output_mode": 'graph', "output_tensor_type": 'padded' }, 'output_mlp': { "use_bias": [True, True, False], "units": [25, 10, 1], "activation": ['relu', 'relu', 'sigmoid'] }, 'attention_args': { "units": 32, 'is_sorted': False, 'has_unconnected': True } } # Update default values input_embedd = update_model_args(model_default['input_embedd'], input_embedd) output_embedd = update_model_args(model_default['output_embedd'], output_embedd) output_mlp = update_model_args(model_default['output_mlp'], output_mlp) attention_args = update_model_args(model_default['attention_args'], attention_args) pooling_nodes_args = {} # Make input embedding, if no feature dimension node_input, n, edge_input, ed, edge_index_input, _, _ = generate_standard_graph_input( input_node_shape, input_edge_shape, None, **input_embedd) edi = edge_index_input nk = Dense(units=attention_args["units"], activation="linear")(n) for i in range(0, depth): heads = [ AttentionHeadGAT(**attention_args)([nk, ed, edi]) for _ in range(attention_heads_num) ] if attention_heads_concat: nk = Concatenate(axis=-1)(heads) else: nk = Average()(heads) n = nk if output_embedd["output_mode"] == 'graph': out = PoolingNodes(**pooling_nodes_args)(n) output_mlp.update({"input_tensor_type": "tensor"}) out = MLP(**output_mlp)(out) main_output = ks.layers.Flatten()(out) # will be dense else: # node embedding out = MLP(**output_mlp)(n) main_output = ChangeTensorType(input_tensor_type="ragged", output_tensor_type="tensor")(out) model = tf.keras.models.Model( inputs=[node_input, edge_input, edge_index_input], outputs=main_output) return model
def make_unet( # Input input_node_shape, input_edge_shape, input_embedd: dict = None, # Output output_embedd: dict = None, output_mlp: dict = None, # Model specific hidden_dim=32, depth=4, k=0.3, score_initializer='ones', use_bias=True, activation='relu', is_sorted=False, has_unconnected=True, use_reconnect=True ): """ Make Graph U Net. Args: input_node_shape (list): Shape of node features. If shape is (None,) embedding layer is used. input_edge_shape (list): Shape of edge features. If shape is (None,) embedding layer is used. input_embedd (list): Dictionary of embedding parameters used if input shape is None. Default is {'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100, 'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64, 'input_type': 'ragged'} output_mlp (dict, optional): Parameter for MLP output classification/ regression. Defaults to {"use_bias": [True, False], "output_dim": [25, 1], "activation": ['relu', 'sigmoid']} output_embedd (str): Dictionary of embedding parameters of the graph network. Default is {"output_mode": 'graph', "output_type": 'padded'} hidden_dim (int): Hidden node feature dimension 32, depth (int): Depth of pooling steps. Default is 4. k (float): Pooling ratio. Default is 0.3. score_initializer (str): How to initialize score kernel. Default is 'ones'. use_bias (bool): Use bias. Default is True. activation (str): Activation function used. Default is 'relu'. is_sorted (bool, optional): Edge edge_indices are sorted. Defaults to False. has_unconnected (bool, optional): Has unconnected nodes. Defaults to True. use_reconnect (bool): Reconnect nodes after pooling. I.e. adj_matrix=adj_matrix^2. Default is True. Returns: model (ks.models.Model): Unet model. """ # Default values update model_default = {'input_embedd': {'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100, 'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64, 'input_tensor_type': 'ragged'}, 'output_embedd': {"output_mode": 'graph', "output_type": 'padded'}, 'output_mlp': {"use_bias": [True, False], "units": [25, 1], "activation": ['relu', 'sigmoid']} } # Update model args input_embedd = update_model_args(model_default['input_embedd'], input_embedd) output_embedd = update_model_args(model_default['output_embedd'], output_embedd) output_mlp = update_model_args(model_default['output_mlp'], output_mlp) pooling_args = {"pooling_method": 'segment_mean', "is_sorted": is_sorted, "has_unconnected": has_unconnected} # Make input embedding, if no feature dimension node_input, n, edge_input, ed, edge_index_input, _, _ = generate_standard_graph_input(input_node_shape, input_edge_shape, None, **input_embedd) tens_type = "values_partition" node_indexing = "batch" n = ChangeTensorType(input_tensor_type="ragged", output_tensor_type=tens_type)(n) ed = ChangeTensorType(input_tensor_type="ragged", output_tensor_type=tens_type)(ed) edi = ChangeTensorType(input_tensor_type="ragged", output_tensor_type=tens_type)(edge_index_input) edi = ChangeIndexing(input_tensor_type=tens_type, to_indexing=node_indexing)([n, edi]) # disjoint output_mlp.update({"input_tensor_type": tens_type}) gather_args = {"input_tensor_type": tens_type, "node_indexing": node_indexing} pooling_args.update({"input_tensor_type": tens_type, "node_indexing": node_indexing}) # Graph lists n = Dense(hidden_dim, use_bias=use_bias, activation='linear', input_tensor_type=tens_type)(n) in_graph = [n, ed, edi] graph_list = [in_graph] map_list = [] # U Down i_graph = in_graph for i in range(0, depth): n, ed, edi = i_graph # GCN layer eu = GatherNodesOutgoing(**gather_args)([n, edi]) eu = Dense(hidden_dim, use_bias=use_bias, activation='linear', input_tensor_type=tens_type)(eu) nu = PoolingLocalEdges(**pooling_args)([n, eu, edi]) # Summing for each node connection n = Activation(activation=activation, input_tensor_type=tens_type)(nu) if use_reconnect: ed, edi = AdjacencyPower(n=2, node_indexing=node_indexing, input_tensor_type=tens_type)([n, ed, edi]) # Pooling i_graph, i_map = PoolingTopK(k=k, kernel_initializer=score_initializer, node_indexing=node_indexing, input_tensor_type=tens_type)([n, ed, edi]) graph_list.append(i_graph) map_list.append(i_map) # U Up ui_graph = i_graph for i in range(depth, 0, -1): o_graph = graph_list[i - 1] i_map = map_list[i - 1] ui_graph = UnPoolingTopK(node_indexing=node_indexing, input_tensor_type=tens_type)(o_graph + i_map + ui_graph) n, ed, edi = ui_graph # skip connection n = Add(input_tensor_type=tens_type)([n, o_graph[0]]) # GCN eu = GatherNodesOutgoing(**gather_args)([n, edi]) eu = Dense(hidden_dim, use_bias=use_bias, activation='linear', input_tensor_type=tens_type)(eu) nu = PoolingLocalEdges(**pooling_args)([n, eu, edi]) # Summing for each node connection n = Activation(activation=activation, input_tensor_type=tens_type)(nu) ui_graph = [n, ed, edi] # Otuput n = ui_graph[0] if output_embedd["output_mode"] == 'graph': out = PoolingNodes(**pooling_args)(n) output_mlp.update({"input_tensor_type": "tensor"}) out = MLP(**output_mlp)(out) main_output = ks.layers.Flatten()(out) # will be dense else: # node embedding out = MLP(**output_mlp)(n) main_output = ChangeTensorType(input_tensor_type=tens_type, output_tensor_type="tensor")(out) model = ks.models.Model(inputs=[node_input, edge_input, edge_index_input], outputs=main_output) return model
def make_nmpn( # Input input_node_shape, input_edge_shape, input_embedd: dict = None, # Output output_embedd: dict = None, output_mlp: dict = None, # Model specific depth=3, node_dim=128, edge_dense: dict = None, use_set2set=True, set2set_args: dict = None, pooling_args: dict = None): """ Get Message passing model. Args: input_node_shape (list): Shape of node features. If shape is (None,) embedding layer is used. input_edge_shape (list): Shape of edge features. If shape is (None,) embedding layer is used. input_embedd (dict): Dictionary of embedding parameters used if input shape is None. Default is {'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100, 'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64, 'input_type': 'ragged'} output_embedd (str): Dictionary of embedding parameters of the graph network. Default is {"output_mode": 'graph', "output_type": 'padded'} output_mlp (dict): Dictionary of MLP arguments for output regression or classifcation. Default is {"use_bias": [True, True, False], "units": [25, 10, 1], "output_activation": ['selu', 'selu', 'sigmoid']} depth (int, optional): Depth. Defaults to 3. node_dim (int, optional): Dimension for hidden node representation. Defaults to 128. edge_dense (dict): Dictionary of arguments for NN to make edge matrix. Default is {'use_bias' : True, 'activation' : 'selu'} use_set2set (bool, optional): Use set2set layer. Defaults to True. set2set_args (dict): Dictionary of Set2Set Layer Arguments. Default is {'channels': 32, 'T': 3, "pooling_method": "sum", "init_qstar": "0"} pooling_args (dict): Dictionary for message pooling arguments. Default is {'is_sorted': False, 'has_unconnected': True, 'pooling_method': "segment_mean"} Returns: model (ks.models.Model): Message Passing model. """ # Make default parameter model_default = { 'input_embedd': { 'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100, 'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64, 'input_tensor_type': 'ragged' }, 'output_embedd': { "output_mode": 'graph', "output_type": 'padded' }, 'output_mlp': { "use_bias": [True, True, False], "units": [25, 10, 1], "activation": ['selu', 'selu', 'sigmoid'] }, 'set2set_args': { 'channels': 32, 'T': 3, "pooling_method": "sum", "init_qstar": "0" }, 'pooling_args': { 'is_sorted': False, 'has_unconnected': True, 'pooling_method': "segment_mean" }, 'edge_dense': { 'use_bias': True, 'activation': 'selu' } } # Update model args input_embedd = update_model_args(model_default['input_embedd'], input_embedd) output_embedd = update_model_args(model_default['output_embedd'], output_embedd) output_mlp = update_model_args(model_default['output_mlp'], output_mlp) set2set_args = update_model_args(model_default['set2set_args'], set2set_args) pooling_args = update_model_args(model_default['pooling_args'], pooling_args) edge_dense = update_model_args(model_default['edge_dense'], edge_dense) # Make input embedding, if no feature dimension node_input, n, edge_input, ed, edge_index_input, _, _ = generate_standard_graph_input( input_node_shape, input_edge_shape, None, **input_embedd) tens_type = "values_partition" node_indexing = "batch" n = ChangeTensorType(input_tensor_type="ragged", output_tensor_type=tens_type)(n) ed = ChangeTensorType(input_tensor_type="ragged", output_tensor_type=tens_type)(ed) edi = ChangeTensorType(input_tensor_type="ragged", output_tensor_type=tens_type)(edge_index_input) edi = ChangeIndexing(input_tensor_type=tens_type, to_indexing=node_indexing)([n, edi]) set2set_args.update({"input_tensor_type": tens_type}) output_mlp.update({"input_tensor_type": tens_type}) edge_dense.update({"input_tensor_type": tens_type}) pooling_args.update({ "input_tensor_type": tens_type, "node_indexing": node_indexing }) n = Dense(node_dim, activation="linear", input_tensor_type=tens_type)(n) edge_net = Dense(node_dim * node_dim, **edge_dense)(ed) gru = GRUupdate(node_dim, input_tensor_type=tens_type, node_indexing=node_indexing) for i in range(0, depth): eu = GatherNodesOutgoing(input_tensor_type=tens_type, node_indexing=node_indexing)([n, edi]) eu = TrafoMatMulMessages(node_dim, input_tensor_type=tens_type, node_indexing=node_indexing)([edge_net, eu]) eu = PoolingLocalEdges(**pooling_args)( [n, eu, edi]) # Summing for each node connections n = gru([n, eu]) if output_embedd["output_mode"] == 'graph': if use_set2set: # output outss = Dense(set2set_args['channels'], activation="linear", input_tensor_type=tens_type)(n) out = Set2Set(**set2set_args)(outss) else: out = PoolingNodes(**pooling_args)(n) # final dense layers output_mlp.update({"input_tensor_type": "tensor"}) main_output = MLP(**output_mlp)(out) else: # Node labeling out = n main_output = MLP(**output_mlp)(out) main_output = ChangeTensorType( input_tensor_type=tens_type, output_tensor_type="tensor")(main_output) # no ragged for distribution supported atm model = ks.models.Model(inputs=[node_input, edge_input, edge_index_input], outputs=main_output) return model
def make_megnet( # Input input_node_shape, input_edge_shape, input_state_shape, input_embedd: dict = None, # Output output_embedd: dict = None, # Only graph possible for megnet output_mlp: dict = None, # Model specs meg_block_args: dict = None, node_ff_args: dict = None, edge_ff_args: dict = None, state_ff_args: dict = None, set2set_args: dict = None, nblocks: int = 3, has_ff: bool = True, dropout: float = None, use_set2set: bool = True, ): """ Get Megnet model. Args: input_node_shape (list): Shape of node features. If shape is (None,) embedding layer is used. input_edge_shape (list): Shape of edge features. If shape is (None,) embedding layer is used. input_state_shape (list): Shape of state features. If shape is (,) embedding layer is used. input_embedd (dict): Dictionary of embedding parameters used if input shape is None. Default is {'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100, 'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64, 'input_type': 'ragged'}. output_embedd (str): Dictionary of embedding parameters of the graph network. Default is {"output_mode": 'graph', "output_type": 'padded'} output_mlp (dict): Dictionary of MLP arguments for output regression or classifcation. Default is {"use_bias": [True, True, True], "units": [32, 16, 1], "activation": ['softplus2', 'softplus2', 'linear']}. meg_block_args (dict): Dictionary of MegBlock arguments. Default is {'node_embed': [64, 32, 32], 'edge_embed': [64, 32, 32], 'env_embed': [64, 32, 32], 'activation': 'softplus2', 'is_sorted': False, 'has_unconnected': True}. node_ff_args (dict): Dictionary of Feed-Forward Layer arguments. Default is {"units": [64, 32], "activation": "softplus2"}. edge_ff_args (dict): Dictionary of Feed-Forward Layer arguments. Default is {"units": [64, 32], "activation": "softplus2"}. state_ff_args (dict): Dictionary of Feed-Forward Layer arguments. Default is {"units": [64, 32], "activation": "softplus2"}. set2set_args (dict): Dictionary of Set2Set Layer Arguments. Default is {'channels': 16, 'T': 3, "pooling_method": "sum", "init_qstar": "0"} nblocks (int): Number of block. Default is 3. has_ff (bool): Use a Feed-Forward layer. Default is True. dropout (float): Use dropout. Default is None. use_set2set (bool): Use set2set. Default is True. Returns: model (tf.keras.models.Model): MEGnet model. """ # Default arguments if None model_default = {'input_embedd': {'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100, 'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64, 'input_tensor_type': 'ragged'}, 'output_embedd': {"output_mode": 'graph', "output_tensor_type": 'padded'}, 'output_mlp': {"use_bias": [True, True, True], "units": [32, 16, 1], "activation": ['kgcnn>softplus2', 'kgcnn>softplus2', 'linear']}, 'meg_block_args': {'node_embed': [64, 32, 32], 'edge_embed': [64, 32, 32], 'env_embed': [64, 32, 32], 'activation': 'kgcnn>softplus2', 'is_sorted': False, 'has_unconnected': True}, 'set2set_args': {'channels': 16, 'T': 3, "pooling_method": "sum", "init_qstar": "0"}, 'node_ff_args': {"units": [64, 32], "activation": "kgcnn>softplus2"}, 'edge_ff_args': {"units": [64, 32], "activation": "kgcnn>softplus2"}, 'state_ff_args': {"units": [64, 32], "activation": "kgcnn>softplus2"} } # Update default arguments input_embedd = update_model_args(model_default['input_embedd'], input_embedd) output_embedd = update_model_args(model_default['output_embedd'], output_embedd) output_mlp = update_model_args(model_default['output_mlp'], output_mlp) meg_block_args = update_model_args(model_default['meg_block_args'], meg_block_args) set2set_args = update_model_args(model_default['set2set_args'], set2set_args) node_ff_args = update_model_args(model_default['node_ff_args'], node_ff_args) edge_ff_args = update_model_args(model_default['edge_ff_args'], edge_ff_args) state_ff_args = update_model_args(model_default['state_ff_args'], state_ff_args) state_ff_args.update({"input_tensor_type": "tensor"}) # Make input embedding, if no feature dimension node_input, n, edge_input, ed, edge_index_input, env_input, uenv = generate_standard_graph_input(input_node_shape, input_edge_shape, input_state_shape, **input_embedd) edi = edge_index_input # starting vp = n ep = ed up = uenv vp = MLP(**node_ff_args)(vp) ep = MLP(**edge_ff_args)(ep) up = MLP(**state_ff_args)(up) vp2 = vp ep2 = ep up2 = up for i in range(0, nblocks): if has_ff and i > 0: vp2 = MLP(**node_ff_args)(vp) ep2 = MLP(**edge_ff_args)(ep) up2 = MLP(**state_ff_args)(up) # MEGnetBlock vp2, ep2, up2 = MEGnetBlock(**meg_block_args)( [vp2, ep2, edi, up2]) # skip connection if dropout is not None: vp2 = Dropout(dropout, name='dropout_atom_%d' % i)(vp2) ep2 = Dropout(dropout, name='dropout_bond_%d' % i)(ep2) up2 = Dropout(dropout, name='dropout_state_%d' % i)(up2) vp = Add()([vp2, vp]) ep = Add()([ep2, ep]) up = Add(input_tensor_type="tensor")([up2, up]) if use_set2set: vp = Dense(set2set_args["channels"], activation='linear')(vp) # to match units ep = Dense(set2set_args["channels"], activation='linear')(ep) # to match units vp = Set2Set(**set2set_args)(vp) ep = Set2Set(**set2set_args)(ep) else: vp = PoolingNodes()(vp) ep = PoolingGlobalEdges()(ep) ep = ks.layers.Flatten()(ep) vp = ks.layers.Flatten()(vp) final_vec = ks.layers.Concatenate(axis=-1)([vp, ep, up]) if dropout is not None: final_vec = ks.layers.Dropout(dropout, name='dropout_final')(final_vec) # final dense layers main_output = MLP(**output_mlp, input_tensor_type="tensor")(final_vec) model = ks.models.Model(inputs=[node_input, edge_input, edge_index_input, env_input], outputs=main_output) return model
def make_schnet( # Input input_node_shape, input_edge_shape, input_embedd: dict = None, # Output output_mlp: dict = None, output_dense: dict = None, output_embedd: dict = None, # Model specific depth=4, out_scale_pos=0, interaction_args: dict = None, node_pooling_args: dict = None): """ Make uncompiled SchNet model. Args: input_node_shape (list): Shape of node features. If shape is (None,) embedding layer is used. input_edge_shape (list): Shape of edge features. If shape is (None,) embedding layer is used. input_embedd (list): Dictionary of embedding parameters used if input shape is None. Default is {'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100, 'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64, 'input_type': 'ragged'} output_mlp (dict, optional): Parameter for MLP output classification/ regression. Defaults to {"use_bias": [True, True], "units": [128, 64], "activation": ['shifted_softplus', 'shifted_softplus']} output_dense (dict): Parameter for Dense scaling layer. Defaults to {"units": 1, "activation": 'linear', "use_bias": True}. output_embedd (str): Dictionary of embedding parameters of the graph network. Default is {"output_mode": 'graph', "output_type": 'padded'} depth (int, optional): Number of Interaction units. Defaults to 4. out_scale_pos (int, optional): Scaling output, position of layer. Defaults to 0. interaction_args (dict): Interaction Layer arguments. Defaults include {"node_dim" : 128, "use_bias": True, "activation" : 'shifted_softplus', "cfconv_pool" : 'segment_sum', "is_sorted": False, "has_unconnected": True} node_pooling_args (dict, optional): Node pooling arguments. Defaults to {"pooling_method": "segment_sum"}. Returns: model (tf.keras.models.Model): SchNet. """ # Make default values if None model_default = { 'input_embedd': { 'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100, 'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64, 'input_tensor_type': 'ragged' }, 'output_embedd': { "output_mode": 'graph', "output_type": 'padded' }, 'interaction_args': { "units": 128, "use_bias": True, "activation": 'shifted_softplus', "cfconv_pool": 'sum', "is_sorted": False, "has_unconnected": True }, 'output_mlp': { "use_bias": [True, True], "units": [128, 64], "activation": ['shifted_softplus', 'shifted_softplus'] }, 'output_dense': { "units": 1, "activation": 'linear', "use_bias": True }, 'node_pooling_args': { "pooling_method": "sum" } } # Update args input_embedd = update_model_args(model_default['input_embedd'], input_embedd) interaction_args = update_model_args(model_default['interaction_args'], interaction_args) output_mlp = update_model_args(model_default['output_mlp'], output_mlp) output_dense = update_model_args(model_default['output_dense'], output_dense) output_embedd = update_model_args(model_default['output_embedd'], output_embedd) node_pooling_args = update_model_args(model_default['node_pooling_args'], node_pooling_args) # Make input embedding, if no feature dimension node_input, n, edge_input, ed, edge_index_input, _, _ = generate_standard_graph_input( input_node_shape, input_edge_shape, None, **input_embedd) # Use representation tens_type = "values_partition" node_indexing = "batch" n = ChangeTensorType(input_tensor_type="ragged", output_tensor_type=tens_type)(n) ed = ChangeTensorType(input_tensor_type="ragged", output_tensor_type=tens_type)(ed) edi = ChangeTensorType(input_tensor_type="ragged", output_tensor_type=tens_type)(edge_index_input) edi = ChangeIndexing(input_tensor_type=tens_type, to_indexing=node_indexing)([n, edi]) n = Dense(interaction_args["units"], activation='linear', input_tensor_type=tens_type)(n) for i in range(0, depth): n = SchNetInteraction(input_tensor_type=tens_type, node_indexing=node_indexing, **interaction_args)([n, ed, edi]) n = MLP(input_tensor_type=tens_type, **output_mlp)(n) mlp_last = Dense(input_tensor_type=tens_type, **output_dense) if output_embedd["output_mode"] == 'graph': if out_scale_pos == 0: n = mlp_last(n) out = PoolingNodes(input_tensor_type=tens_type, node_indexing=node_indexing, **node_pooling_args)(n) if out_scale_pos == 1: out = mlp_last(out) main_output = ks.layers.Flatten()(out) # will be dense else: # node embedding out = mlp_last(n) main_output = ChangeTensorType( input_tensor_type="values_partition", output_tensor_type="tensor")(out) # no ragged for distribution atm model = ks.models.Model(inputs=[node_input, edge_input, edge_index_input], outputs=main_output) return model
def make_gcn( # Input input_node_shape, input_edge_shape, input_embedd: dict = None, # Output output_embedd: dict = None, output_mlp: dict = None, # Model specific depth=3, gcn_args: dict = None): """ Make GCN model. Args: input_node_shape (list): Shape of node features. If shape is (None,) embedding layer is used. input_edge_shape (list): Shape of edge features. If shape is (None,) embedding layer is used. input_embedd (dict): Dictionary of embedding parameters used if input shape is None. Default is {"input_node_vocab": 100, "input_edge_vocab": 10, "input_state_vocab": 100, "input_node_embedd": 64, "input_edge_embedd": 64, "input_state_embedd": 64, "input_type": 'ragged'}. output_embedd (dict): Dictionary of embedding parameters of the graph network. Default is {"output_mode": 'graph', "output_type": 'padded'}. output_mlp (dict): Dictionary of arguments for final MLP regression or classifcation layer. Default is {"use_bias": [True, True, False], "units": [25, 10, 1], "activation": ['relu', 'relu', 'sigmoid']}. depth (int, optional): Number of convolutions. Defaults to 3. gcn_args (dict): Dictionary of arguments for the GCN convolutional unit. Defaults to {"units": 100, "use_bias": True, "activation": 'relu', "pooling_method": 'segment_sum', "is_sorted": False, "has_unconnected": "True"}. Returns: model (tf.keras.models.Model): uncompiled model. """ if input_edge_shape[-1] != 1: raise ValueError( "No edge features available for GCN, only edge weights of pre-scaled adjacency matrix, \ must be shape (batch, None, 1), but got (without batch-dimension): ", input_edge_shape) # Make default args model_default = { 'input_embedd': { "input_node_vocab": 100, "input_edge_vocab": 10, "input_state_vocab": 100, "input_node_embedd": 64, "input_edge_embedd": 64, "input_state_embedd": 64, "input_tensor_type": 'ragged' }, 'output_embedd': { "output_mode": 'graph', "output_tensor_type": 'masked' }, 'output_mlp': { "use_bias": [True, True, False], "units": [25, 10, 1], "activation": ['relu', 'relu', 'sigmoid'] }, 'gcn_args': { "units": 100, "use_bias": True, "activation": 'relu', "pooling_method": 'sum', "is_sorted": False, "has_unconnected": True } } # Update model parameter input_embedd = update_model_args(model_default['input_embedd'], input_embedd) output_embedd = update_model_args(model_default['output_embedd'], output_embedd) output_mlp = update_model_args(model_default['output_mlp'], output_mlp) gcn_args = update_model_args(model_default['gcn_args'], gcn_args) # Make input embedding, if no feature dimension node_input, n, edge_input, ed, edge_index_input, env_input, uenv = generate_standard_graph_input( input_node_shape, input_edge_shape, None, **input_embedd) edi = edge_index_input # Map to units n = Dense(gcn_args["units"], use_bias=True, activation='linear')(n) # n-Layer Step for i in range(0, depth): n = GCN(**gcn_args)([n, ed, edi]) if output_embedd["output_mode"] == "graph": out = PoolingNodes()(n) # will return tensor output_mlp.update({"input_tensor_type": "tensor"}) out = MLP(**output_mlp)(out) else: # Node labeling out = n out = MLP(**output_mlp)(out) out = ChangeTensorType( input_tensor_type='ragged', output_tensor_type="tensor")( out) # no ragged for distribution supported atm model = ks.models.Model(inputs=[node_input, edge_input, edge_index_input], outputs=out) return model
def make_attentiveFP( # Input input_node_shape, input_edge_shape, input_embedd: dict = None, # Output output_embedd: dict = None, output_mlp: dict = None, # Model specific parameter depth=3, attention_args: dict = None): """ Generate Interaction network. Args: input_node_shape (list): Shape of node features. If shape is (None,) embedding layer is used. input_edge_shape (list): Shape of edge features. If shape is (None,) embedding layer is used. input_embedd (dict): Dictionary of embedding parameters used if input shape is None. Default is {'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100, 'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64, 'input_type': 'ragged'}. output_embedd (dict): Dictionary of embedding parameters of the graph network. Default is {"output_mode": 'graph', "output_type": 'padded'}. output_mlp (dict): Dictionary of arguments for final MLP regression or classifcation layer. Default is {"use_bias": [True, True, False], "units": [25, 10, 1], "activation": ['relu', 'relu', 'sigmoid']}. depth (int): Number of convolution layers. Default is 3. attention_args (dict): Layer arguments for attention layer. Default is {"units": 32, 'is_sorted': False, 'has_unconnected': True} Returns: model (tf.keras.model): Interaction model. """ print("Warning model has not been tested yet.") # default values model_default = { 'input_embedd': { 'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100, 'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64, 'input_tensor_type': 'ragged' }, 'output_embedd': { "output_mode": 'graph', "output_tensor_type": 'padded' }, 'output_mlp': { "use_bias": [True, True, False], "units": [25, 10, 1], "activation": ['relu', 'relu', 'sigmoid'] }, 'attention_args': { "units": 32, 'is_sorted': False, 'has_unconnected': True } } # Update default values input_embedd = update_model_args(model_default['input_embedd'], input_embedd) output_embedd = update_model_args(model_default['output_embedd'], output_embedd) output_mlp = update_model_args(model_default['output_mlp'], output_mlp) attention_args = update_model_args(model_default['attention_args'], attention_args) # Make input embedding, if no feature dimension node_input, n, edge_input, ed, edge_index_input, _, _ = generate_standard_graph_input( input_node_shape, input_edge_shape, None, **input_embedd) edi = edge_index_input nk = Dense(units=attention_args['units'])(n) Ck = AttentiveHeadFP(use_edge_features=True, **attention_args)([nk, ed, edi]) nk = GRUupdate(units=attention_args['units'])([nk, Ck]) for i in range(1, depth): Ck = AttentiveHeadFP(**attention_args)([nk, ed, edi]) nk = GRUupdate(units=attention_args['units'])([nk, Ck]) n = nk if output_embedd["output_mode"] == 'graph': out = AttentiveNodePooling(units=attention_args['units'])(n) output_mlp.update({"input_tensor_type": "tensor"}) out = MLP(**output_mlp)(out) main_output = ks.layers.Flatten()(out) # will be dense else: # node embedding out = MLP(**output_mlp)(n) main_output = ChangeTensorType(input_tensor_type="ragged", output_tensor_type="tensor")(out) model = tf.keras.models.Model( inputs=[node_input, edge_input, edge_index_input], outputs=main_output) return model
def make_dimnet_pp( # Input input_node_shape, input_embedd: dict = None, # Output output_embedd: dict = None, # Model specific parameter emb_size = 128, out_emb_size = 256, int_emb_size = 64, basis_emb_size =8, num_blocks = 4, num_spherical = 7, num_radial= 6, cutoff=5.0, envelope_exponent=5, num_before_skip=1, num_after_skip=2, num_dense_output=3, num_targets=12, activation="swish", extensive=True, output_init='zeros', ): model_default = {'input_embedd': {'input_node_vocab': 95, 'input_node_embedd': 64, 'input_tensor_type': 'ragged'} } input_embedd = update_model_args(model_default['input_embedd'], input_embedd) node_input, n, xyz_input, bond_index_input, angle_index_input, _ = generate_mol_graph_input(input_node_shape, [None, 3], [None, 2], [None, 2], **input_embedd) x = xyz_input edi = bond_index_input adi = angle_index_input # Calculate distances d = NodeDistance()([x, edi]) rbf = BesselBasisLayer(num_radial=num_radial, cutoff=cutoff, envelope_exponent=envelope_exponent)(d) # Calculate angles a = EdgeAngle()([x, edi, adi]) sbf = SphericalBasisLayer(num_spherical=num_spherical, num_radial=num_radial, cutoff=cutoff, envelope_exponent=envelope_exponent)([d, a, adi]) # Embedding block rbf_emb = Dense(emb_size, use_bias=True, activation=activation, kernel_initializer="orthogonal")(rbf) n_pairs = GatherNodes()([n, edi]) x = Concatenate(axis=-1)([n_pairs, rbf_emb]) x = Dense(emb_size, use_bias=True, activation=activation, kernel_initializer="orthogonal")(x) ps = DimNetOutputBlock(emb_size, out_emb_size, num_dense_output, num_targets=num_targets, output_kernel_initializer=output_init)([n, x, rbf, edi]) # Interaction blocks add_xp = Add() for i in range(num_blocks): x = DimNetInteractionPPBlock(emb_size, int_emb_size, basis_emb_size, num_before_skip, num_after_skip)( [x, rbf, sbf, adi]) p_update = DimNetOutputBlock(emb_size, out_emb_size, num_dense_output, num_targets=num_targets, output_kernel_initializer=output_init)([n, x, rbf, edi]) ps = add_xp([ps, p_update]) if extensive: main_output = PoolingNodes(pooling_method="sum")(ps) else: main_output = PoolingNodes(pooling_method="mean")(ps) model = tf.keras.models.Model(inputs=[node_input, xyz_input, bond_index_input, angle_index_input], outputs=main_output) return model
def make_graph_sage( # Input input_node_shape, input_edge_shape, input_embedd: dict = None, # Output output_embedd: dict = None, output_mlp: dict = None, # Model specific parameter depth=3, use_edge_features=False, node_mlp_args: dict = None, edge_mlp_args: dict = None, pooling_args: dict = None): """ Generate Interaction network. Args: input_node_shape (list): Shape of node features. If shape is (None,) embedding layer is used. input_edge_shape (list): Shape of edge features. If shape is (None,) embedding layer is used. input_embedd (dict): Dictionary of embedding parameters used if input shape is None. Default is {'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100, 'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64, 'input_type': 'ragged'}. output_embedd (dict): Dictionary of embedding parameters of the graph network. Default is {"output_mode": 'graph', "output_type": 'padded'}. output_mlp (dict): Dictionary of arguments for final MLP regression or classifcation layer. Default is {"use_bias": [True, True, False], "units": [25, 10, 1], "activation": ['relu', 'relu', 'sigmoid']}. depth (int): Number of convolution layers. Default is 3. use_edge_features (bool): Whether to concatenate edges with nodes in aggregate. Default is False. node_mlp_args (dict): Dictionary of arguments for MLP for node update. Default is {"units": [100, 50], "use_bias": True, "activation": ['relu', "linear"]} edge_mlp_args (dict): Dictionary of arguments for MLP for interaction update. Default is {"units": [100, 100, 100, 100, 50], "activation": ['relu', 'relu', 'relu', 'relu', "linear"]} pooling_args (dict): Dictionary for message pooling arguments. Default is {'is_sorted': False, 'has_unconnected': True, 'pooling_method': "segment_mean"} Returns: model (tf.keras.model): Interaction model. """ # default values model_default = { 'input_embedd': { 'input_node_vocab': 95, 'input_edge_vocab': 5, 'input_state_vocab': 100, 'input_node_embedd': 64, 'input_edge_embedd': 64, 'input_state_embedd': 64, 'input_tensor_type': 'ragged' }, 'output_embedd': { "output_mode": 'graph', "output_tensor_type": 'padded' }, 'output_mlp': { "use_bias": [True, True, False], "units": [25, 10, 1], "activation": ['relu', 'relu', 'sigmoid'] }, 'node_mlp_args': { "units": [100, 50], "use_bias": True, "activation": ['relu', "linear"] }, 'edge_mlp_args': { "units": [100, 50], "use_bias": True, "activation": ['relu', "linear"] }, 'pooling_args': { 'is_sorted': False, 'has_unconnected': True, 'pooling_method': "segment_mean" } } # Update default values input_embedd = update_model_args(model_default['input_embedd'], input_embedd) output_embedd = update_model_args(model_default['output_embedd'], output_embedd) output_mlp = update_model_args(model_default['output_mlp'], output_mlp) node_mlp_args = update_model_args(model_default['node_mlp_args'], node_mlp_args) edge_mlp_args = update_model_args(model_default['edge_mlp_args'], edge_mlp_args) pooling_args = update_model_args(model_default['pooling_args'], pooling_args) pooling_nodes_args = { "input_tensor_type": 'ragged', "node_indexing": 'sample', 'pooling_method': "mean" } gather_args = {"node_indexing": 'sample'} concat_args = {"axis": -1, "input_tensor_type": 'ragged'} # Make input embedding, if no feature dimension node_input, n, edge_input, ed, edge_index_input, _, _ = generate_standard_graph_input( input_node_shape, input_edge_shape, None, **input_embedd) edi = edge_index_input for i in range(0, depth): # upd = GatherNodes()([n,edi]) eu = GatherNodesOutgoing(**gather_args)([n, edi]) if use_edge_features: eu = Concatenate(**concat_args)([eu, ed]) eu = MLP(**edge_mlp_args)(eu) # Pool message if pooling_args['pooling_method'] in ["LSTM", "lstm"]: nu = PoolingLocalEdgesLSTM(**pooling_args)([n, eu, edi]) else: nu = PoolingLocalMessages(**pooling_args)( [n, eu, edi]) # Summing for each node connection nu = Concatenate(**concat_args)( [n, nu]) # Concatenate node features with new edge updates n = MLP(**node_mlp_args)(nu) n = LayerNormalization(axis=-1)(n) # Normalize # Regression layer on output if output_embedd["output_mode"] == 'graph': out = PoolingNodes(**pooling_nodes_args)(n) output_mlp.update({"input_tensor_type": "tensor"}) out = MLP(**output_mlp)(out) main_output = ks.layers.Flatten()(out) # will be tensor else: # node embedding out = MLP(**output_mlp)(n) main_output = ChangeTensorType(input_tensor_type='ragged', output_tensor_type="tensor")(out) model = tf.keras.models.Model( inputs=[node_input, edge_input, edge_index_input], outputs=main_output) return model