91 lines
3.0 KiB
Python
91 lines
3.0 KiB
Python
from federatedscope.core.configs.config import CN
|
|
from federatedscope.register import register_config
|
|
|
|
|
|
def extend_model_cfg(cfg):
|
|
# ---------------------------------------------------------------------- #
|
|
# Model related options
|
|
# ---------------------------------------------------------------------- #
|
|
cfg.model = CN()
|
|
|
|
cfg.model.model_num_per_trainer = 1 # some methods may leverage more
|
|
# than one model in each trainer
|
|
cfg.model.type = 'lr'
|
|
cfg.model.use_bias = True
|
|
cfg.model.task = 'node'
|
|
cfg.model.hidden = 256
|
|
cfg.model.dropout = 0.5
|
|
cfg.model.in_channels = 0 # If 0, model will be built by data.shape
|
|
cfg.model.out_channels = 1
|
|
cfg.model.layer = 2 # In GPR-GNN, K = layer
|
|
cfg.model.graph_pooling = 'mean'
|
|
cfg.model.embed_size = 8
|
|
cfg.model.num_item = 0
|
|
cfg.model.num_user = 0
|
|
cfg.model.input_shape = () # A tuple, e.g., (in_channel, h, w)
|
|
|
|
# For tree-based model
|
|
cfg.model.lambda_ = 0.1
|
|
cfg.model.gamma = 0
|
|
cfg.model.num_of_trees = 10
|
|
cfg.model.max_tree_depth = 3
|
|
|
|
# language model for hetero NLP tasks
|
|
cfg.model.stage = '' # ['assign', 'contrast']
|
|
cfg.model.model_type = 'google/bert_uncased_L-2_H-128_A-2'
|
|
cfg.model.pretrain_tasks = []
|
|
cfg.model.downstream_tasks = []
|
|
cfg.model.num_labels = 1
|
|
cfg.model.max_length = 200
|
|
cfg.model.min_length = 1
|
|
cfg.model.no_repeat_ngram_size = 3
|
|
cfg.model.length_penalty = 2.0
|
|
cfg.model.num_beams = 5
|
|
cfg.model.label_smoothing = 0.1
|
|
cfg.model.n_best_size = 20
|
|
cfg.model.max_answer_len = 30
|
|
cfg.model.null_score_diff_threshold = 0.0
|
|
cfg.model.use_contrastive_loss = False
|
|
cfg.model.contrast_topk = 100
|
|
cfg.model.contrast_temp = 1.0
|
|
|
|
# Traffic Flow model parameters, These are only default values.
|
|
# Please modify the specific parameters directly in the baselines/YAML files.
|
|
cfg.model.num_nodes = 0
|
|
cfg.model.rnn_units = 64
|
|
cfg.model.dropout = 0.1
|
|
cfg.model.horizon = 12
|
|
cfg.model.input_dim = 1 # If 0, model will be built by data.shape
|
|
cfg.model.output_dim = 1
|
|
cfg.model.embed_dim = 10
|
|
cfg.model.num_layers = 1 # In GPR-GNN, K = layer
|
|
cfg.model.cheb_order = 1 # A tuple, e.g., (in_channel, h, w)
|
|
cfg.model.use_day = True
|
|
cfg.model.use_week = True
|
|
|
|
|
|
# ---------------------------------------------------------------------- #
|
|
# Criterion related options
|
|
# ---------------------------------------------------------------------- #
|
|
cfg.criterion = CN()
|
|
|
|
cfg.criterion.type = 'MSELoss'
|
|
|
|
# ---------------------------------------------------------------------- #
|
|
# regularizer related options
|
|
# ---------------------------------------------------------------------- #
|
|
cfg.regularizer = CN()
|
|
|
|
cfg.regularizer.type = ''
|
|
cfg.regularizer.mu = 0.
|
|
|
|
# --------------- register corresponding check function ----------
|
|
cfg.register_cfg_check_fun(assert_model_cfg)
|
|
|
|
|
|
def assert_model_cfg(cfg):
|
|
pass
|
|
|
|
|
|
register_config("model", extend_model_cfg)
|