pipeline.GNN.models package

A package that defines various triplet-based GNNs.

pipeline.GNN.models.get_model(model_type=None)

Get a GNN model class from its name.

Parameters:

model_type (Optional[str]) – GNN type. Available are triplet_interaction, edge_based and scifi_triplet_interaction.

Return type:

Type[TripletGNNBase]

Returns:

The GNN model class that can be instantiated.

pipeline.GNN.models.edge_based_gnn module

class pipeline.GNN.models.edge_based_gnn.EdgeBasedGNN(hparams)

Bases: TripletGNNBase

A triplet-based GNN without using node encodings.

forward_edges(x, start, end)

Forward step for edge classification.

Parameters:

• x (Tensor) – Hit features

• edge_index – Torch tensor with 2 rows that define the edges.

Returns:

the hit encodings and edge encodings after message passing, and the edge classifier output.

Return type:

A tuple of 3 tensors

forward_triplets(e, filtered_edge_index, edge_mask, dict_triplet_indices, **kwargs)

Forward step for triplet classification.

Parameters:

• e (Tensor) – Edge encodings after the edge forward step

• filtered_edge_index (Tensor) – edge index after requiring the minimal edge score

• edge_mask (Tensor) – edge mask of edges kept after filtering

• edge_score_cut – custom edge score cut. If not specified, the value is taken from the internal hparams dictionary.

Return type:

Dict[str, Tensor]

Returns:

A dictionary that associates articulation, elbow_left and elbow_right with the logits of the corresponding triplets.

property input_kwargs: Dict[str, Any]

Associates an input name with a dictionary corresponding to the keyword arguments used to build a dummy tensor representing the input. This dictionary basically gives the size and dtype of the tensor.

property input_to_dynamic_axes

A dictionary that associates an input name with the dynamic axis specification.

message_step(e, start, end, dim_size)

Return type:

Tensor

property subnetwork_groups: Dict[str, List[str]]

A dictionary that associates a subnetwork actually corresponding to a list of subnetworks, with this list of subnetworks.

property subnetwork_to_outputs: Dict[str, List[str]]

A dictionary that associates a subnetwork name with the list of its output names.

property subnetworks: List[str]

List of subnetworks available. It is derived from subnetwork_to_outputs.

triplet_output_step_articulation(e, edge_indices, triplet_indices)

triplet_output_step_elbow_left(e, edge_indices, triplet_indices)

triplet_output_step_elbow_right(e, edge_indices, triplet_indices)

pipeline.GNN.models.incremental_triplet_interaction_gnn module

pipeline.GNN.models.scifi_triplet_interaction_gnn module

pipeline.GNN.models.triplet_interaction_gnn module

class pipeline.GNN.models.triplet_interaction_gnn.TripletInteractionGNN(*args: Any, **kwargs: Any)

Bases: TripletGNNBase

A triplet-based interaction network.

encoding_step(x, start, end)

Initial encoding step of the GNN.

Parameters:

• x (Tensor) – hit input features

• start (Tensor) – start indices of the edges

• end (Tensor) – end indices of the edges

Return type:

Tuple[Tensor, Tensor]

Returns:

The node encodings h and edge encodings e

forward_edges(x, start, end)

Forwrd step for edge classification.

Parameters:

• x (Tensor) – hit input features

• start (Tensor) – start indices of the edges

• end (Tensor) – end indices of the edges

Returns:

the hit encodings and edge encodings after message passing, and the edge classifier output.

Return type:

A tuple of 3 tensors

forward_triplets(h, e, filtered_edge_index, edge_mask, dict_triplet_indices, **kwargs)

Forward step for triplet classification.

Parameters:

• h (Tensor) – Hit encodings after the edge forward step

• e (Tensor) – Edge encodings after the edge forward step

• filtered_edge_index (Tensor) – edge index after requiring the minimal edge score

• edge_mask (Tensor) – edge mask of edges kept after filtering

• edge_score_cut – custom edge score cut. If not specified, the value is taken from the internal hparams dictionary.

Return type:

Dict[str, Tensor]

Returns:

A dictionary that associates articulation, elbow_left and elbow_right with the logits of the corresponding triplets.

property input_kwargs: Dict[str, Any]

Associates an input name with a dictionary corresponding to the keyword arguments used to build a dummy tensor representing the input. This dictionary basically gives the size and dtype of the tensor.

property input_to_dynamic_axes

A dictionary that associates an input name with the dynamic axis specification.

message_step(h, start, end, e, step)

Apply one step of message-passing that updates the node and edge encodings.

Return type:

Tuple[Tensor, Tensor]

property n_graph_iters: int

Number of message-passing iterations

property only_e: bool

output_step(h, start, end, e)

Apply the edge output classifier to edges to get edge logits.

Return type:

Tensor

property recursive: bool

Whether the GNN is recursive, i.e., the networks involved in the message step are the same.

scatter_add(source, index, h)

Scatter add operation. In ONNX export mode for TensorRT, the operation is replaced by a fake operator that is implemented through a plugin in TensorRT.

Return type:

Tensor

scatter_max(source, index, h)

property subnetwork_groups: Dict[str, List[str]]

A dictionary that associates a subnetwork actually corresponding to a list of subnetworks, with this list of subnetworks.

property subnetwork_to_outputs: Dict[str, List[str]]

A dictionary that associates a subnetwork name with the list of its output names.

to_onnx(outpath, mode=None, options=None)

Export the model to ONNX

Parameters:

• outpath (str) – path to the ONNX output file

• mode (Optional[str]) – subnetwork to save

Return type:

None

triplet_output_step_articulation(h, e, edge_indices, triplet_indices)

triplet_output_step_elbow_left(h, e, edge_indices, triplet_indices)

triplet_output_step_elbow_right(h, e, edge_indices, triplet_indices)