openGauss数据库源码解析系列文章——AI技术（2.1）

def main(args):
logging.basicConfig(level=logging.INFO)
# 实例化算法模型，模板化模型或DNN模型
model = SQLDiag(args.model, args.csv_file, get_config(args.config_file))
# 训练模型
if args.mode == 'train':
    # fit训练数据，提取模板或特征
        model.fit()
        # 模型保存
        model.save(args.model_path)
    # 预测
elif args.mode == 'predict':
    # 加载模型
        model.load(args.model_path)
        # 标准化预测数据，获取结果
        pred_result = model.transform()
        # 保存输出结果
        ResultSave().save(pred_result, args.predicted_file)
        logging.info('predict result in saved in {}'.format(args.predicted_file))
    # 更新模型
elif args.mode == 'finetune':
        model.fine_tune(args.model_path)
        model.save(args.model_path)

class TemplateModel(AbstractModel):
    # 初始化算法参数
    def __init__(self, params):
        super().__init__(params)
        self.bias = 1e-5
        self.__hash_table = dict(INSERT=dict(), UPDATE=dict(), DELETE=dict(), SELECT=dict(),
                                 OTHER=dict())
        self.time_list_size = params.time_list_size
        self.knn_number = params.knn_number
        self.similarity_algorithm = calc_sql_distance(params.similarity_algorithm)

def fit(self, data):
    # 对每条sql语句按照粗、细粒度进行标准化，生成模板
        for sql, duration_time in data:
            if not self.check_illegal_sql(sql):
                continue
            fine_template, rough_template = get_sql_template(sql)
            sql_prefix = fine_template.split()[0]
            if sql_prefix not in self.__hash_table:
                sql_prefix = 'OTHER'
            # 更新粗粒度模板框架
            if rough_template not in self.__hash_table[sql_prefix]:
                self.__hash_table[sql_prefix][rough_template] = dict()
                self.__hash_table[sql_prefix][rough_template]['info'] = dict()
            # 更新细粒度模板框架
            if fine_template not in self.__hash_table[sql_prefix][rough_template]['info']:
                self.__hash_table[sql_prefix][rough_template]['info'][fine_template] = \
                    dict(time_list=[], count=0, mean_time=0.0, iter_time=0.0)
            # 更新每个细粒度模板的执行时间、迭代时间、sql语句的计数。
            …

self.__hash_table[sql_prefix][rough_template]['info'][fine_template]['count'] += 1
…        
# 基于细粒度模板更新粗粒度模板信息
        for sql_prefix, sql_prefix_info in self.__hash_table.items():
            …

    def transform(self, data):
        predict_time_list = {}
        for sql in data:
            # sql语句不属于'INSERT', 'SELECT', 'UPDATE', 'DELETE', 'CREATE', 'DROP'任何一个，预测时间默认为-1
            if not self.check_illegal_sql(sql):
                predict_time_list[sql] = -1
                continue
            …
                # 若预测的sql所对应的粗粒度模板不存在，执行模板相似度计算方法获取与所有粗粒度模板的相似度
                if rough_template not in self.__hash_table[sql_prefix]:
                    for local_rough_template, local_rough_template_info in self.__hash_table[
                            sql_prefix].items():
                        similarity_info.append(
                            (self.similarity_algorithm(rough_template, local_rough_template), local_rough_template_info['mean_time']))
                # 若预测的sql所对应的细粒度模板不存在，执行模板相似度计算方法获取与所有细粒度模板的相似度
                else:
                    for local_fine_template, local_fine_template_info in \
                            self.__hash_table[sql_prefix][rough_template][
                                'info'].items():
                        similarity_info.append(
                            (self.similarity_algorithm(fine_template, local_fine_template),
                             local_fine_template_info['iter_time']))
                # 基于KNN思想计算sql执行时间
                topn_similarity_info = heapq.nlargest(self.knn_number, similarity_info)
                …

        return predict_time_list

class KerasRegression:
    # 初始化模型参数
    def __init__(self, encoding_dim=1):
        self.model = None
        self.encoding_dim = encoding_dim
# 模型定义
    @staticmethod
    def build_model(shape, encoding_dim):
        from tensorflow.keras import Input, Model
        from tensorflow.keras.layers import Dense
        inputs = Input(shape=(shape,))
        layer_dense1 = Dense(128, activatinotallow='relu', kernel_initializer='he_normal')(inputs)
        …
        model = Model(inputs=inputs, outputs=y_pred)
        # 优化器，损失函数
        model.compile(optimizer='adam', loss='mse', metrics=['mae'])
        return model
    # 模型训练
    def fit(self, features, labels, batch_size=128, epochs=300):
        …
        self.model.fit(features, labels, epochs=epochs, batch_size=batch_size, shuffle=True, verbose=2)
    # 模型预测
    def predict(self, features):
        predict_result = self.model.predict(features)
        return predict_result
    # 模型保存
    def save(self, filepath):
        self.model.save(filepath)
    # 模型读取
    def load(self, filepath):
        from tensorflow.keras.models import load_model
        self.model = load_model(filepath)

class DnnModel(AbstractModel, ABC):
    # 初始化算法参数
    def __init__(self, params):
        …
        self.regression = KerasRegression(encoding_dim=1)
        self.data = None
    # 把sql语句转化为vector，如果模型不存在，则直接训练w2v模型，如果模型存在则进行增量训练
    def build_word2vector(self, data):
        self.data = list(data)
        if self.w2v.model:
            self.w2v.update(self.data)
        else:
            self.w2v.fit(self.data)

    def fit(self, data):
        self.build_word2vector(data)
        …
        # 数据归一化
        self.scaler = MinMaxScaler(feature_range=(0, 1))
        self.scaler.fit(labels)
        labels = self.scaler.transform(labels)
        self.regression.fit(features, labels, epochs=self.epoch)

# 利用回归模型预测执行时间
def transform(self, data):
…

python main.py –help

python main.py train -f train.csv --model-path test/

python main.py predict –f predict.csv –model-path test/ --predicted-file test/result.csv

python main.py finetune –f train_new.csv –model-path test/