小白也想搞科研（二）之代码升级

import pandas as pd
import numpy as np
import random
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
torch.autograd.set_detect_anomaly(True)


# 生成模拟数据库查询日志的函数
def generate_query_log(num_queries=1000):
    fields = ["id", "name", "age", "email"]
    conditions = ["= ?", "> ?", "< ?", "LIKE ?"]

    query_log = pd.DataFrame({
        "query_id": range(1, num_queries + 1),
        "num_fields": [random.randint(1, len(fields)) for _ in range(num_queries)],
        "num_conditions": [random.randint(1, 4) for _ in range(num_queries)],
        "use_index": [random.choice([True, False]) for _ in range(num_queries)],
        "data_volume": [random.randint(100, 10000) for _ in range(num_queries)]  # 模拟的数据量
    })

    # 计算资源消耗
    base_resource_usage_per_field = 0.05 # 字段数对资源消耗的影响较小
    base_resource_usage_per_condition = 0.15 # 条件数对资源消耗的影响较大
    base_resource_usage_per_volume_unit = 0.001  # 每单位数据量的资源消耗
    query_log["resource_usage"] = query_log["num_fields"] * base_resource_usage_per_field + query_log["num_conditions"] * base_resource_usage_per_condition + query_log["data_volume"] * base_resource_usage_per_volume_unit

    # 设置执行时间
    base_execution_time = 2.0 # 设定一个基础执行时间
    execution_time_factor = 0.5 # 用于调整执行时间与资源消耗关系的因子
    query_log["execution_time"] = base_execution_time + (query_log["resource_usage"] * execution_time_factor)

    query_log["query"] = query_log.apply(lambda row: f"SELECT {', '.join(random.sample(fields, int(row['num_fields'])))} FROM users WHERE {' AND '.join([random.choice(fields) + ' ' + random.choice(conditions) for _ in range(int(row['num_conditions']))])}", axis=1)

    return query_log




def extract_features(row):
    # 提取查询的长度、涉及的字段数、条件数、是否使用索引、以及数据量
    query_length = len(row['query'])
    num_fields = row['num_fields']
    num_conditions = row['num_conditions']
    use_index = 1 if row['use_index'] else 0  # 将布尔值转换为整数
    data_volume = row['data_volume']  # 查询涉及的数据量

    # 将这些特征组合成一个特征向量
    features = [query_length, num_fields, num_conditions, use_index, data_volume]

    return features


# 定义策略网络
# 策略网络是一个神经网络，它从状态空间接收输入，并输出每个动作的概率。
# 在我们的例子中，状态空间简化为查询的长度和其中空格的数量（作为查询复杂性的简化代理）
# 而动作空间是一个二元选择：使用索引或不使用索引。

#定义了一个新的类 PolicyNetwork，它继承自 PyTorch 的 nn.Module。
# 在 PyTorch 中，nn.Module 是所有神经网络模型的基类。
# 定义的 PolicyNetwork 是一个特定的神经网络，用于实现策略学习。
class PolicyNetwork(nn.Module):
    #初始化函数。它定义了网络的结构。该函数接收两个参数：num_inputs（输入层的维度，即状态空间的大小）
    # 和 num_actions（输出层的维度，即动作空间的大小）。
    def __init__(self, num_inputs, num_actions):
        #调用父类 nn.Module 的初始化函数。在创建新的 PyTorch 模型时，这是标准的做法，它设置了一些背后的基础设施。
        # 增加了网络的容量，并添加了一个额外的隐藏层来处理更复杂的输入。这样可以帮助网络学习更复杂的模式，并更好地理解数据库查询的状态。
        super(PolicyNetwork, self).__init__()
        self.fc1 = nn.Linear(num_inputs, 128)  # 定义了网络的第一层，是一个全连接（线性）层。创建了一个从 num_inputs 维到 128 维的线性映射。128 是这一层神经元的数量。
        self.fc2 = nn.Linear(128, 64)  # 可以添加一个额外的隐藏层
        self.fc3 = nn.Linear(64, num_actions)  # 第三层，这一层从第二层的 64 个神经元接收输入，并将其映射到 num_actions 个输出，每个输出对应一个动作的概率。
#定义了数据通过网络的前向传播过程，每当你对 PolicyNetwork 的实例进行调用时，这个函数就会被执行。
    def forward(self, x):
        #在前向传播中，数据首先通过第一层，然后应用ReLU激活函数。ReLU（Rectified Linear Unit）是一个非线性函数，它将所有负值置为0，对正值不做改变。这有助于引入非线性，使网络能够学习更复杂的模式。
        x = torch.relu(self.fc1(x))  # 激活函数ReLU
        x = torch.relu(self.fc2(x))  # 第二层ReLU激活
        #最后，数据通过第三层，并应用softmax函数。softmax函数可以将原始输出转换为概率分布，这些概率总和为1。在这里，它给出了选择每个动作的概率。
        return torch.softmax(self.fc3(x), dim=1)  # Softmax输出动作概率


# 定义奖励函数
def reward_function(execution_time, resource_usage, max_time, max_resource_usage):
    # 标准化执行时间和资源消耗
    normalized_time = execution_time / max_time
    normalized_resource = resource_usage / max_resource_usage

    time_weight = 0.7
    resource_weight = 0.3

    # 调整奖励计算公式
    reward = (1 / normalized_time) * time_weight - normalized_resource * resource_weight
    return reward

# 奖励函数是根据执行时间来计算奖励的。在这个简化模型中，奖励是执行时间的倒数，意味着更快的执行时间会获得更高的奖励。


# 在现实情况中，查询的复杂性对于是否应该使用索引有重大影响。一般来说，涉及大量数据的复杂查询可能从索引中受益更多。
def simulate_index_usage(execution_time, resource_usage):
    # 假设使用索引会减少执行时间
    time_reduction_factor = 0.5  # 可以根据实际情况调整
    # 假设使用索引会增加资源消耗
    resource_usage_increase_factor = 1.2  # 可以根据实际情况调整

    adjusted_execution_time = execution_time * (1 - time_reduction_factor)
    adjusted_resource_usage = resource_usage * resource_usage_increase_factor

    return adjusted_execution_time, adjusted_resource_usage

def simulate_without_index(execution_time, resource_usage):
    # 不使用索引时，执行时间和资源消耗保持不变
    return execution_time, resource_usage



# 训练模型
def train(policy_net, optimizer, query_log, num_episodes=1000):
    max_execution_time = query_log['execution_time'].max()  # 获取最大执行时间用于标准化
    max_resource_usage = query_log['resource_usage'].max()  # 获取最大资源消耗用于标准化
    epsilon = 0.1  # 探索率
    for episode in range(num_episodes):
        total_reward = 0
        for _, row in query_log.iterrows():
            # 更新状态表示，包含所有新特征
            state = torch.tensor([extract_features(row)], dtype=torch.float32)
            if random.random() < epsilon:  # 探索
                action = torch.tensor([random.choice([0, 1])], dtype=torch.int)
            else:  # 利用
                action_probs = policy_net(state)
                m = Categorical(action_probs)
                action = m.sample()

            # 根据当前的策略网络选择一个动作，然后根据执行时间计算奖励
            # execution_time = row['execution_time'] if action.item() == 0 else row['execution_time'] / 2
            if action.item() == 0:  # 假设动作0代表使用索引
                execution_time, resource_usage = simulate_index_usage(row['execution_time'], row['resource_usage'])
            else:  # 不使用索引
                execution_time, resource_usage = simulate_without_index(row['execution_time'], row['resource_usage'])

            reward = reward_function(execution_time, resource_usage, max_execution_time, max_resource_usage)
            total_reward += reward

            # 使用这个奖励更新策略网络，目的是让网络学习如何选择可以获得更高奖励的动作。
            optimizer.zero_grad()
            # 重新计算log_prob来避免共享计算图
            log_prob = Categorical(policy_net(state)).log_prob(action)
            loss = -log_prob * reward
            loss.backward()
            optimizer.step()
        print(f'Episode {episode} Total Reward: {total_reward}')

# 生成模拟数据
query_log = generate_query_log()

# 设置DRL环境和模型
num_features = 5  # 这里是状态空间的维度
num_actions = 2  # 二元动作空间：使用索引或不使用索引

policy_net = PolicyNetwork(num_features, num_actions)
optimizer = optim.Adam(policy_net.parameters(), lr=0.01)

# 训练模型
train(policy_net, optimizer, query_log)