graphnet.models.gnn.icemix module

Implementation of IceMix architecture used in.

IceCube - Neutrinos in Deep Ice

Reconstruct the direction of neutrinos from the Universe to the South Pole

Kaggle competition.

Solution by DrHB: https://github.com/DrHB/icecube-2nd-place

class graphnet.models.gnn.icemix.DeepIce(*args, **kwargs)[source]

Bases: GNN

DeepIce model.

Construct DeepIce.

Parameters:

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

  • mlp_ratio (int, default: 4) – Mlp expansion ratio of FourierEncoder and Transformer.

  • seq_length (int, default: 192) – The base feature dimension.

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

  • head_size (int, default: 32) – The size of the attention heads.

  • depth_rel (int, default: 4) – The depth of the relative transformer.

  • n_rel (int, default: 1) – The number of relative transformer layers to use.

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

  • include_dynedge (bool, default: False) – If True, pulse-level predictions from DynEdge will be added as features to the model.

  • dynedge_args (Optional[Dict[str, Any]], default: None) – Initialization arguments for DynEdge. If not provided, DynEdge will be initialized with the original Kaggle Competition settings. If include_dynedge is False, this argument have no impact.

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

  • args (Any)

  • kwargs (Any)

Return type:

object

no_weight_decay()[source]

cls_tocken should not be subject to weight decay during training.

Return type:

Set

forward(data)[source]

Apply learnable forward pass.

Return type:

Tensor

Parameters:

data (Data)