KTeleBERT/config.py

import os.path as osp
import numpy as np
import random
import torch
from easydict import EasyDict as edict
import argparse


LAYER_MAPPING = {
    0: 'od_layer_0',
    1: 'od_layer_1',
    2: 'od_layer_2',
}


class cfg():
    def __init__(self):
        self.this_dir = osp.dirname(__file__)
        # change
        self.data_root = osp.abspath(osp.join(self.this_dir, '..', '..', 'data', ''))

        # TODO: add some static variable  (The frequency of change is low)

    def get_args(self):
        parser = argparse.ArgumentParser()
        # ------------ base ------------
        parser.add_argument('--train_strategy', default=1, type=int)
        parser.add_argument('--batch_size', default=64, type=int)
        parser.add_argument('--batch_size_ke', default=14, type=int)
        parser.add_argument('--batch_size_od', default=8, type=int)
        parser.add_argument('--batch_size_ad', default=32, type=int)

        parser.add_argument('--epoch', default=15, type=int)
        parser.add_argument("--save_model", default=1, type=int, choices=[0, 1])
        # 用transformer的 save_pretrain 方式保存
        parser.add_argument("--save_pretrain", default=0, type=int, choices=[0, 1])
        parser.add_argument("--from_pretrain", default=0, type=int, choices=[0, 1])

        # torthlight
        parser.add_argument("--no_tensorboard", default=False, action="store_true")
        parser.add_argument("--exp_name", default="huawei_exp", type=str, help="Experiment name")
        parser.add_argument("--dump_path", default="dump/", type=str, help="Experiment dump path")
        parser.add_argument("--exp_id", default="ke256_raekt_ernie2_bs20_p3_c3_5e-6", type=str, help="Experiment ID")
        # or 3407
        parser.add_argument("--random_seed", default=42, type=int)
        # 数据参数
        parser.add_argument("--data_path", default="huawei", type=str, help="Experiment path")
        parser.add_argument('--train_ratio', default=1, type=float, help='ratio for train/test')
        parser.add_argument("--seq_data_name", default='Seq_data_base', type=str, help="seq_data 名字")
        parser.add_argument("--kg_data_name", default='KG_data_base_rule', type=str, help="kg_data 名字")
        parser.add_argument("--order_data_name", default='event_order_data', type=str, help="order_data 名字")
        # TODO: add some dynamic variable
        parser.add_argument("--model_name", default="MacBert", type=str, help="model name")

        # ------------ 训练阶段 ------------
        parser.add_argument("--scheduler", default="cos", type=str, choices=["linear", "cos"])
        parser.add_argument("--optim", default="adamw", type=str)
        parser.add_argument("--adam_epsilon", default=1e-8, type=float)
        parser.add_argument('--workers', type=int, default=8)
        parser.add_argument('--accumulation_steps', type=int, default=6)
        parser.add_argument('--accumulation_steps_ke', type=int, default=6)
        parser.add_argument('--accumulation_steps_ad', type=int, default=6)
        parser.add_argument('--accumulation_steps_od', type=int, default=6)
        parser.add_argument("--train_together", default=0, type=int)

        # 3e-5
        parser.add_argument('--lr', type=float, default=1e-5)
        # 逐层学习率衰减
        parser.add_argument("--LLRD", default=0, type=int, choices=[0, 1])
        parser.add_argument('--weight_decay', type=float, default=0.01)
        parser.add_argument('--clip', type=float, default=1., help='gradient clipping')
        parser.add_argument('--scheduler_steps', type=int, default=None,
                            help='total number of step for the scheduler, if None then scheduler_total_step = total_step')
        parser.add_argument('--eval_step', default=100, type=int, help='evaluate each n step')

        # ------------ PLM ------------
        parser.add_argument('--maxlength', type=int, default=200)
        parser.add_argument('--mlm_probability', type=float, default=0.15)
        parser.add_argument('--final_mlm_probability', type=float, default=0.4)
        parser.add_argument('--mlm_probability_increase', type=str, default="curve", choices=["linear", "curve"])
        parser.add_argument("--mask_stratege", default="rand", type=str, choices=["rand", "wwm", "domain"])
        # 前n个epoch 用rand，后面用wwm. multi-stage knowledge masking strategy
        parser.add_argument("--ernie_stratege", default=-1, type=int)
        # 用mlm任务进行训练,默认使用chinese_ref且添加新的special word
        parser.add_argument("--use_mlm_task", default=1, type=int, choices=[0, 1])
        # 添加新的special word
        parser.add_argument("--add_special_word", default=1, type=int, choices=[0, 1])
        # freeze
        parser.add_argument("--freeze_layer", default=0, type=int, choices=[0, 1, 2, 3, 4])
        # 是否mask 特殊token
        parser.add_argument("--special_token_mask", default=0, type=int, choices=[0, 1])
        parser.add_argument("--emb_init", default=1, type=int, choices=[0, 1])
        parser.add_argument("--cls_head_init", default=1, type=int, choices=[0, 1])
        # 是否使用自适应权重
        parser.add_argument("--use_awl", default=1, type=int, choices=[0, 1])
        parser.add_argument("--mask_loss_scale", default=1.0, type=float)

        # ------------ KGE ------------
        parser.add_argument('--ke_norm', type=int, default=1)
        parser.add_argument('--ke_dim', type=int, default=768)
        parser.add_argument('--ke_margin', type=float, default=1.0)
        parser.add_argument('--neg_num', type=int, default=10)
        parser.add_argument('--adv_temp', type=float, default=1.0, help='The temperature of sampling in self-adversarial negative sampling.')
        # 5e-4
        parser.add_argument('--ke_lr', type=float, default=3e-5)
        parser.add_argument('--only_ke_loss', type=int, default=0)

        # ------------ 数值embedding相关 ------------
        parser.add_argument('--use_NumEmb', type=int, default=1)
        parser.add_argument("--contrastive_loss", default=1, type=int, choices=[0, 1])
        parser.add_argument("--l_layers", default=2, type=int)
        parser.add_argument('--use_kpi_loss', type=int, default=1)

        # ------------ 测试阶段 ------------
        parser.add_argument("--only_test", default=0, type=int, choices=[0, 1])
        parser.add_argument("--mask_test", default=0, type=int, choices=[0, 1])
        parser.add_argument("--embed_gen", default=0, type=int, choices=[0, 1])
        parser.add_argument("--ke_test", default=0, type=int, choices=[0, 1])
        # -1: 测全集
        parser.add_argument("--ke_test_num", default=-1, type=int)
        parser.add_argument("--path_gen", default="", type=str)

        # ------------ 时序阶段 ------------
        # 1：预训练
        # 2：时序 finetune
        # 3. 异常检测 finetune + 时序, 且是迭代的
        # 是否加载od模型
        parser.add_argument("--order_load", default=0, type=int)
        parser.add_argument("--order_num", default=2, type=int)
        parser.add_argument("--od_type", default='linear_cat', type=str, choices=['linear_cat', 'vertical_attention'])
        parser.add_argument("--eps", default=0.2, type=float, help='label smoothing..')
        parser.add_argument("--num_od_layer", default=0, type=int)
        parser.add_argument("--plm_emb_type", default='cls', type=str, choices=['cls', 'last_avg'])
        parser.add_argument("--order_test_name", default='', type=str)
        parser.add_argument("--order_threshold", default=0.5, type=float)
        # ------------ 并行训练 ------------
        # 是否并行
        parser.add_argument('--rank', type=int, default=0, help='rank to dist')
        parser.add_argument('--dist', type=int, default=0, help='whether to dist')
        # 不要改该参数，系统会自动分配
        parser.add_argument('--device', default='cuda', help='device id (i.e. 0 or 0,1 or cpu)')
        # 开启的进程数(注意不是线程),不用设置该参数，会根据nproc_per_node自动设置
        parser.add_argument('--world-size', default=4, type=int,
                            help='number of distributed processes')
        parser.add_argument('--dist-url', default='env://', help='url used to set up distributed training')
        parser.add_argument("--local_rank", default=-1, type=int)
        self.cfg = parser.parse_args()

    def update_train_configs(self):
        # add some constraint for parameters
        # e.g. cannot save and test at the same time
        # 修正默认参数
        # TODO: 测试逻辑有问题需要修改
        if len(self.cfg.order_test_name) > 0:
            self.cfg.save_model = 0
            if len(self.cfg.order_test_name) == 0:
                self.cfg.train_ratio = min(0.8, self.cfg.train_ratio)
            # 自适应载入文件名
            else:
                print("od test ... ")
                self.cfg.train_strategy == 5
                self.cfg.plm_emb_type = 'last_avg' if 'last_avg' in self.cfg.model_name else 'cls'
                for key in LAYER_MAPPING.keys():
                    if LAYER_MAPPING[key] in self.cfg.model_name:
                        self.cfg.num_od_layer = key
                self.cfg.order_test_name = osp.join('downstream_task', f'{self.cfg.order_test_name}')

        if self.cfg.mask_test or self.cfg.embed_gen or self.cfg.ke_test or len(self.cfg.order_test_name) > 0:
            assert len(self.cfg.model_name) > 0
            self.cfg.only_test = 1
        if self.cfg.only_test == 1:
            self.save_model = 0
            self.save_pretrain = 0

        # TODO: update some dynamic variable
        self.cfg.data_root = self.data_root
        self.cfg.data_path = osp.join(self.data_root, self.cfg.data_path)
        self.cfg.plm_path = osp.join(self.data_root, 'transformer')
        self.cfg.dump_path = osp.join(self.cfg.data_path, self.cfg.dump_path)
        # bs 控制尽量在32

        # 自适应权重的数量
        self.cfg.awl_num = 1
        # ------------ 数值embedding相关 ------------
        self.cfg.hidden_size = 768
        self.cfg.num_attention_heads = 8
        self.cfg.hidden_dropout_prob = 0.1
        self.cfg.num_kpi = 304
        self.cfg.specail_emb_path = None
        if self.cfg.emb_init:
            self.cfg.specail_emb_path = osp.join(self.cfg.data_path, 'added_vocab_embedding.pt')

        # ------------- 多任务学习相关 -------------
        # 四个阶段
        self.cfg.mask_epoch, self.cfg.ke_epoch, self.cfg.ad_epoch, self.cfg.od_epoch = None, None, None, None
        # 触发多任务 学习
        if self.cfg.train_strategy > 1:
            self.cfg.mask_epoch = [0, 1, 1, 1, 0]
            self.cfg.ke_epoch = [4, 3, 2, 2, 0]
            if self.cfg.only_ke_loss:
                self.cfg.mask_epoch = [0, 0, 0, 0, 0]
            self.cfg.epoch = sum(self.cfg.mask_epoch) + sum(self.cfg.ke_epoch)
            if self.cfg.train_strategy > 2:
                self.cfg.ad_epoch = [0, 6, 3, 1, 0]
                self.cfg.epoch += sum(self.cfg.ad_epoch)
                if self.cfg.train_strategy > 3 and not self.cfg.only_ke_loss:
                    self.cfg.od_epoch = [0, 0, 9, 1, 0]
                    # self.cfg.mask_epoch[3] = 1
                    self.cfg.epoch += sum(self.cfg.od_epoch)
            self.cfg.epoch_matrix = []
            for epochs in [self.cfg.mask_epoch, self.cfg.ke_epoch, self.cfg.ad_epoch, self.cfg.od_epoch]:
                if epochs is not None:
                    self.cfg.epoch_matrix.append(epochs)
            if self.cfg.train_together:
                # loss 直接相加，训练epoch就是mask的epoch
                self.cfg.epoch = sum(self.cfg.mask_epoch)
                self.cfg.batch_size = int((self.cfg.batch_size - 16) / self.cfg.train_strategy)
                self.cfg.batch_size_ke = int(self.cfg.batch_size_ke / self.cfg.train_strategy) - 2
                self.cfg.batch_size_ad = int(self.cfg.batch_size_ad / self.cfg.train_strategy) - 1
                self.cfg.batch_size_od = int(self.cfg.batch_size_od / self.cfg.train_strategy) - 1
                self.cfg.accumulation_steps = (self.cfg.accumulation_steps - 1) * self.cfg.train_strategy

        self.cfg.neg_num = max(min(self.cfg.neg_num, self.cfg.batch_size_ke - 3), 1)

        self.cfg.accumulation_steps_dict = {0: self.cfg.accumulation_steps, 1: self.cfg.accumulation_steps_ke, 2: self.cfg.accumulation_steps_ad, 3: self.cfg.accumulation_steps_od}

        # 使用数值embedding也必须添加新词因为位置信息和tokenizer绑定
        if self.cfg.use_mlm_task or self.cfg.use_NumEmb:
            assert self.cfg.add_special_word == 1

        if self.cfg.use_NumEmb:
            self.cfg.awl_num += 1

        return self.cfg