graphnet.models.transformer.iseecube module

Implementation of ISeeCube Transformer architecture used in.

https://github.com/ChenLi2049/ISeeCube/

class graphnet.models.transformer.iseecube.ISeeCube(*args, **kwargs)[source]

Bases: GNN

ISeeCube model.

Construct ISeeCube.

Parameters:

  • hidden_dim (int, default: 384) – The latent feature dimension.

  • seq_length (int, default: 196) – The number of pulses in a neutrino event.

  • num_layers (int, default: 16) – The depth of the transformer.

  • num_heads (int, default: 12) – The number of the attention heads.

  • mlp_dim (int, default: 1536) – The mlp dimension of FourierEncoder and Transformer.

  • rel_pos_buckets (int, default: 32) – Relative position buckets for relative position bias.

  • max_rel_pos (int, default: 256) – Maximum relative position for relative position bias.

  • num_register_tokens (int, default: 3) – The number of register tokens.

  • scaled_emb (bool, default: False) – Whether to scale the sinusoidal positional embeddings.

  • n_features (int, default: 6) – The number of features in the input data.

  • args (Any)

  • kwargs (Any)

Return type:

object

forward(data)[source]

Apply learnable forward pass.

Return type:

Tensor

Parameters:

data (Data)