lloca.backbone.graphnet.GraphNet

class lloca.backbone.graphnet.GraphNet(in_channels, hidden_reps, out_channels, num_blocks, *args, checkpoint_blocks=False, **kwargs)

Bases: Module

Baseline LLoCa-GNN.

Simple message-passing graph neural network, consisting of EdgeConv blocks.

Parameters:

• in_channels (int) – Number of input channels.

• hidden_reps (str) – Tensor representation used in the hidden layers.

• out_channels (int) – Number of output channels.

• num_blocks (int) – Number of EdgeConv blocks.

• *args

• checkpoint_blocks (bool) – Whether to use gradient checkpointing in the EdgeConv blocks.

• **kwargs

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(inputs, frames, edge_index, batch=None, edge_attr=None)

Forward pass.

Parameters:

• inputs (Tensor) – Input data with shape (num_items, in_channels)

• frames (Frames) – Local frames used for message passing

• edge_index (Tensor) – Edge index tensor with shape (2, num_edges)

• batch (Tensor) – Batch tensorwith shape (num_items,) If None, assumes fully connected graph along the num_items direction.

• edge_attr (Tensor) – Edge attribute tensor with shape (num_edges, num_edge_attr)

Returns:

outputs – Outputs with shape (num_items, out_channels)

Return type:

Tensor