torchdrug.tasks

Property Prediction Tasks

PropertyPrediction

class PropertyPrediction(model, task=(), criterion='mse', metric='mae', 'rmse', verbose=0)

Graph / molecule property prediction task.

This class is also compatible with semi-supervised learning.

Parameters

• model (nn.Module) – graph representation model

• task (str, list or dict, optional) – training task(s). For dict, the keys are tasks and the values are the corresponding weights.

• criterion (str, list or dict, optional) – training criterion(s). For dict, the keys are criterions and the values are the corresponding weights. Available criterions are mse and bce.

• metric (str or list of str, optional) – metric(s). Available metrics are mae, rmse, auprc and auroc.

• verbose (int, optional) – output verbose level

preprocess(train_set, valid_set, test_set)

Compute the mean and derivation for each task on the training set.

Pre-trained Molecular Representation Tasks

EdgePrediction

class EdgePrediction(model)

Edge prediction task proposed in Inductive Representation Learning on Large Graphs.

Parameters

model (nn.Module) – node representation model

AttributeMasking

class AttributeMasking(model, mask_rate=0.15, num_mlp_layer=2)

Attribute masking proposed in Strategies for Pre-training Graph Neural Networks.

Parameters

• model (nn.Module) – node representation model

• mask_rate (float, optional) – rate of masked nodes

• num_mlp_layer (int, optional) – number of MLP layers

ContextPrediction

class ContextPrediction(model, context_model=None, k=5, r1=4, r2=7, readout='mean', num_negative=1)

Context prediction task proposed in Strategies for Pre-training Graph Neural Networks.

For a given center node, the subgraph is defined as a k-hop neighborhood (inclusive) around the selected node. The context graph is defined as the surrounding graph structure between r1- (exclusive) and r2-hop (inclusive) from the center node. Nodes between k- and r1-hop are picked as anchor nodes for the context representation.

Parameters

• model (nn.Module) – node representation model for subgraphs.

• context_model (nn.Module, optional) – node representation model for context graphs. By default, use the same architecture as model without parameter sharing.

• k (int, optional) – radius for subgraphs

• r1 (int, optional) – inner radius for context graphs

• r2 (int, optional) – outer radius for context graphs

• readout (nn.Module, optional) – readout function over context anchor nodes

• num_negative (int, optional) – number of negative samples per positive sample

Unsupervised

class Unsupervised(model)

Wrapper task for unsupervised learning.

The unsupervised loss should be computed by the model.

Parameters

model (nn.Module) – any model

Molecule Generation Tasks

AutoregressiveGeneration

class AutoregressiveGeneration(node_model, edge_model, task=(), num_node_sample=- 1, num_edge_sample=- 1, max_edge_unroll=None, max_node=None, criterion='nll', agent_update_interval=5, gamma=0.9, reward_temperature=1, baseline_momentum=0.9)

Autoregressive graph generation task.

This class can be used to implement GraphAF proposed in GraphAF: A Flow-based Autoregressive Model for Molecular Graph Generation. To do so, instantiate the node model and the edge model with two GraphAutoregressiveFlow models.

Parameters

• node_model (nn.Module) – node likelihood model

• edge_model (nn.Module) – edge likelihood model

• task (str or list of str, optional) – property optimization task(s). Available tasks are plogp and qed.

• num_node_sample (int, optional) – number of node samples per graph. -1 for all samples.

• num_edge_sample (int, optional) – number of edge samples per graph. -1 for all samples.

• max_edge_unroll (int, optional) – max node id difference. If not provided, use the statistics from the training set.

• max_node (int, optional) – max number of node. If not provided, use the statistics from the training set.

• criterion (str, list or dict, optional) – training criterion(s). For dict, the keys are criterions and the values are the corresponding weights. Available criterions are nll and ppo.

• agent_update_interval (int, optional) – update agent every n batch

• gamma (float, optional) – reward discount rate

• reward_temperature (float, optional) – temperature for reward. Higher temperature encourages larger mean reward, while lower temperature encourages larger maximal reward.

• baseline_momentum (float, optional) – momentum for value function baseline

preprocess(train_set, valid_set, test_set)

Add atom id mapping and random BFS order to the training set.

Compute max_edge_unroll and max_node on the training set if not provided.

GCPNGeneration

class GCPNGeneration(model, atom_types, max_edge_unroll=None, max_node=None, task=(), criterion='nll', hidden_dim_mlp=128, agent_update_interval=10, gamma=0.9, reward_temperature=1, baseline_momentum=0.9)

The graph generative model from Graph Convolutional Policy Network for Goal-Directed Molecular Graph Generation.

Parameters

• model (nn.Module) – graph representation model

• atom_types (list or set) – set of all possible atom types

• task (str or list of str, optional) – property optimization task(s)

• max_edge_unroll (int, optional) – max node id difference. If not provided, use the statistics from the training set.

• max_node (int, optional) – max number of node. If not provided, use the statistics from the training set.

• criterion (str, list or dict, optional) – training criterion(s). For dict, the keys are criterions and the values are the corresponding weights. Available criterions are nll and ppo.

• agent_update_interval (int, optional) – update the agent every n batch

• gamma (float, optional) – reward discount rate

• reward_temperature (float, optional) – temperature for reward. Higher temperature encourages larger mean reward, while lower temperature encourages larger maximal reward.

• baseline_momentum (float, optional) – momentum for value function baseline

preprocess(train_set, valid_set, test_set)

Add atom id mapping and random BFS order to the training set.

Compute max_edge_unroll and max_node on the training set if not provided.

Retrosynthesis Tasks

CenterIdentification

class CenterIdentification(model, feature='reaction', 'graph', 'atom', 'bond', num_mlp_layer=2)

Reaction center identification task.

This class is a part of retrosynthesis prediction.

Parameters

• model (nn.Module) – graph representation model

• feature (str or list of str, optional) – additional features for prediction. Available features are reaction: type of the reaction graph: graph representation of the product atom: original atom feature bond: original bond feature

• num_mlp_layer (int, optional) – number of MLP layers

predict_synthon(batch, k=1)

Predict top-k synthons from target molecules.

Parameters

• batch (dict) – batch of target molecules

• k (int, optional) – return top-k results

Returns

top k records.

Each record is a batch dict of keys synthon, num_synthon, reaction_center, log_likelihood and reaction.

Return type

list of dict

SynthonCompletion

class SynthonCompletion(model, feature='reaction', 'graph', 'atom', num_mlp_layer=2)

Synthon completion task.

This class is a part of retrosynthesis prediction.

Parameters

• model (nn.Module) – graph representation model

• feature (str or list of str, optional) – additional features for prediction. Available features are reaction: type of the reaction graph: graph representation of the synthon atom: original atom feature

• num_mlp_layer (int, optional) – number of MLP layers

Retrosynthesis

class Retrosynthesis(center_identification, synthon_completion, center_topk=2, num_synthon_beam=10, max_prediction=20, metric='top-1', 'top-3', 'top-5', 'top-10')

Retrosynthesis task.

This class wraps pretrained center identification and synthon completion modeules into a pipeline.

Parameters

• center_identification (CenterIdentification) – sub task of center identification

• synthon_completion (SynthonCompletion) – sub task of synthon completion

• center_topk (int, optional) – number of reaction centers to predict for each product

• num_synthon_beam (int, optional) – size of beam search for each synthon

• max_prediction (int, optional) – max number of final predictions for each product

• metric (str or list of str, optional) – metric(s). Available metrics are top-K.

load_state_dict(state_dict, strict=True)

Copies parameters and buffers from state_dict into this module and its descendants. If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Parameters

• state_dict (dict) – a dict containing parameters and persistent buffers.

• strict (bool, optional) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True

Returns

• missing_keys is a list of str containing the missing keys

• unexpected_keys is a list of str containing the unexpected keys

Return type

NamedTuple with missing_keys and unexpected_keys fields

Knowledge Graph Reasoning Tasks

KnowledgeGraphCompletion

class KnowledgeGraphCompletion(model, criterion='bce', metric='mr', 'mrr', 'hits@1', 'hits@3', 'hits@10', num_negative=128, margin=6, adversarial_temperature=0, strict_negative=True, filtered_ranking=True, fact_ratio=None, sample_weight=True)

Knowledge graph completion task.

This class provides routines for the family of knowledge graph embedding models.

Parameters

• model (nn.Module) – knowledge graph embedding model

• criterion (str, list or dict, optional) – training criterion(s). For dict, the keys are criterions and the values are the corresponding weights. Available criterions are bce, ce and ranking.

• metric (str or list of str, optional) – metric(s). Available metrics are mr, mrr and hits@K.

• num_negative (int, optional) – number of negative samples per positive sample

• margin (float, optional) – margin in ranking criterion

• adversarial_temperature (float, optional) – temperature for self-adversarial negative sampling. Set 0 to disable self-adversarial negative sampling.

• strict_negative (bool, optional) – use strict negative sampling or not

• filtered_ranking (bool, optional) – use filtered or unfiltered ranking for evaluation

• fact_ratio (float, optional) – split the training set into facts and labels. Set None to use the whole training set as both facts and labels.

• sample_weight (bool, optional) – whether to down-weight triplets from entities of large degrees