parameter.py

import argparse
import json
import yaml

def get_args():
    parser = argparse.ArgumentParser()

    ### Data
    parser.add_argument("--name", type=str, default="Cora", choices=['Cora', 'CiteSeer', 'PubMed', 'Photo', 'Computers'])
    parser.add_argument("--data_path_origin", type=str, default="../Dataset/origin/")
    parser.add_argument("--data_path", type=str, default="../Dataset/gpm_attack/")
    parser.add_argument("--device", type=str, default="cuda:0")
    parser.add_argument("--test_idx_num", type=int, default=100)
    parser.add_argument("--unify_dim", type=int, default=50)
    
 
    ### Train
    parser.add_argument("--gfm_model", type=str, default="GCOPE")
    parser.add_argument("--backdoor_model", type=str, default="ours")
    parser.add_argument("--pretrain", type=int, default=1)
    parser.add_argument("--backdoor", type=int, default=1)
    parser.add_argument("--retrain_backdoor", type=int, default=1)
    parser.add_argument("--backbone_tuning", type=int, default=0)
    
    ### Trigger
    parser.add_argument("--num_trigger_node", type=int, default=3)

    ### Log
    parser.add_argument("--index", type=str, default="")

    args = parser.parse_args()

    with open(f'./config/gfm/{args.gfm_model}.yaml', 'r') as f:
        config1 = yaml.safe_load(f)
    
    with open(f'./config/bd/{args.backdoor_model}.yaml', 'r') as f:
        config2 = yaml.safe_load(f)

    merge_config = {**config1, **config2}
    merge_args = {**merge_config, **vars(args)}
    args = argparse.Namespace(**merge_args)

    all_datasets = ['Cora', 'CiteSeer', 'PubMed', 'Photo', 'Computers']
    # all_datasets = ['Cora', 'CiteSeer', 'Photo']
    args.pretrain_dataset = [d for d in all_datasets if d != args.name]

    return args