config_pretrain.yaml

gpu: 'cuda:0'
lr: 2e-4
min_lr: 1e-7
weight_decay: 0.0
epochs: 5
warmup_epochs: 0.7
patience_epochs: 0.3
load_model: None
log_every_n_steps: 50

# Select the GNN model. Supported models:
# - SchNet:  K. T. Schütt et al., https://aip.scitation.org/doi/10.1063/1.5019779 
# - SE(3)-Transformer: F. B. Fuchs et al., https://arxiv.org/abs/2006.10503
# - EGNN: V. G. Satorras et al., https://arxiv.org/abs/2102.09844 
# - TorchMD-Net: P. Thölke et al., https://arxiv.org/abs/2202.02541
# The defualt settings are listed below.
model: 
  # name: "SchNet"
  # num_atoms: 28
  # bond_feat_dim: 1
  # num_targets: 1
  # hidden_channels: 256
  # num_filters: 256
  # num_interactions: 5
  # num_gaussians: 32
  # cutoff: 5.0
  # max_num_neighbors: 32
  # readout: 'add'

  # name: 'SE3Transformer'
  # num_layers: 3
  # atom_feature_size: 28
  # num_channels: 8
  # num_degrees: 4
  # edge_dim: 4
  # div: 4
  # pooling: 'avg'
  # n_heads: 2
  # cutoff: 5.0
  # max_num_neighbors: 32

  name: "EGNN"
  hidden_channels: 256
  in_edge_nf: 0
  n_layers: 5
  residual: True
  attention: True
  normalize: True
  tanh: False
  cutoff: 5.0
  max_atom_type: 28
  max_chirality_type: 5
  max_num_neighbors: 32

  # name: "TorchMD-Net"
  # hidden_channels: 256
  # num_layers: 6
  # num_rbf: 32
  # rbf_type: "expnorm"
  # trainable_rbf: True
  # activation: "silu"
  # attn_activation: "silu"
  # neighbor_embedding: True
  # num_heads: 8
  # distance_influence: "both"
  # cutoff_lower: 0.0
  # cutoff_upper: 5.0
  # max_atom_type: 28
  # max_chirality_type: 5
  # max_num_neighbors: 32

dataset:
  batch_size: 256
  num_workers: 8
  valid_size: 0.05
  ani1: True
  ani1x: True
  std: 0.2
  seed: 777