tabtransformer.yaml

name: TABTransformer
framework: torch
model_type: tabular

model_tasks:
  - classification

label_types:
  - binary
  - categorical

set_params_function:
  _target_: CMC_utils.models.set_tabtransformer_params

init_params:
  _target_: CMC_utils.models.TABTransformer
  _convert_: all

  cat_idxs:

  categories:

  num_continuous:

  dim: 32

  dim_out:

  depth: 6

  heads: 8

  attn_dropout: 0.1

  ff_dropout: 0.1

  mlp_hidden_mults: [4, 2]

  mlp_act:
    _target_: torch.nn.ReLU

  # device: ${dl_params.device}
  # continuous_mean_std:

fit_params: {}

train_function:
  _target_: CMC_utils.models.train_torch_model

test_function:
  _target_: CMC_utils.models.test_torch_model

save_function:
  _target_: CMC_utils.save_load.save_model

file_extension: pth

load_function:
  _target_: CMC_utils.save_load.load_model