run_model.py

"""
训练并评估单一模型的脚本
"""

import argparse

from libcity.pipeline import run_model
from libcity.utils import str2bool, add_general_args


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # 增加指定的参数
    parser.add_argument('--task', type=str,
                        default='traffic_state_pred', help='the name of task')
    parser.add_argument('--model', type=str,
                        default='AGCRN', help='the name of model')
    parser.add_argument('--dataset', type=str,
                        default='xunfei', help='the name of dataset')
    parser.add_argument('--config_file', type=str,
                        default=None, help='the file name of config file')
    parser.add_argument('--saved_model', type=str2bool,
                        default=True, help='whether save the trained model')
    parser.add_argument('--train', type=str2bool, default=True,
                        help='whether re-train model if the model is trained before')
    parser.add_argument('--exp_id', type=str, default=None, help='id of experiment')
    parser.add_argument('--seed', type=int, default=0, help='random seed')
    parser.add_argument('--K', type=int, default=0, help='K-fold')
    parser.add_argument('--ind', type=int, default=-1, help='select index in K-fold')
    parser.add_argument("--loss", type=str, default='mse')
    parser.add_argument("--rnn_units", type=int, default=128)
    parser.add_argument("--embed_dim", type=int, default=30)
    parser.add_argument("--cheb_order", type=int, default=4)
    parser.add_argument("--num_layers", type=int, default=2)
    parser.add_argument("--graph_type", type=str, default='geo')
    parser.add_argument("--dtw_topk", type=int, default=10)
    parser.add_argument("--add_apt", type=str2bool, default=True)
    parser.add_argument("--use_cl", type=str2bool, default=False)
    parser.add_argument("--step_size", type=int, default=5000)
    parser.add_argument("--path", type=str, default=None)
    # 增加其他可选的参数
    add_general_args(parser)
    # 解析参数
    args = parser.parse_args()
    dict_args = vars(args)
    other_args = {key: val for key, val in dict_args.items() if key not in [
        'task', 'model', 'dataset', 'config_file', 'saved_model', 'train'] and
        val is not None}
    run_model(task=args.task, model_name=args.model, dataset_name=args.dataset,
              config_file=args.config_file, saved_model=args.saved_model,
              train=args.train, other_args=other_args)