CN117234710A - 一种采用强化学习实现ai模型训练内存优化的方法 - Google Patents

Abstract

本发明针对神经网络训练内存优化问题，设计了一种采用强化学习实现AI模型训练内存优化的方法。该方法为神经网络训练过程中产生的中间特征数据搜索分配合适的内存优化策略，方法基于神经网络模型结构确定内存优化策略空间，构建用于模拟神经网络训练过程的多维状态矩阵；采用双智能体强化学习方法协同进行策略搜索；根据策略搜索结果的奖励值，多轮迭代优化以得到收益最大的策略搜索结果，在达到计算加速设备内存优化目标的同时，避免密集化的通信和计算资源消耗；根据策略搜索结果修改神经网络模型的图结构，减少模型训练过程中GPU内存消耗，从而使得单个GPU设备可以处理更大规模的神经网络模型或者更大batch‑size的数据样本。

Description

一种采用强化学习实现AI模型训练内存优化的方法

技术领域

一种采用强化学习实现AI模型训练内存优化的方法，用于减少神经网络训练过程中的GPU内存消耗。

技术背景

神经网络已经被证明具有显著的可伸缩性，通常可以用更多的数据样本来训练更大的模型的方式提高模型精度。但是计算加速设备(如GPU)的内存容量增长速度远小于神经网络训练内存需求的增长速度。如何减少模型在训练过程中的内存消耗成为影响深度学习发展的重要问题。

神经网络模型的训练过程通常包括正向传播和反向传播两个主要阶段。在正常传播过程中提取数据的特征信息，反向传播过程中计算模型参数的误差信息。模型中包含多个计算结点，结点与结点之间具有数据依赖性，上一个结点的输出会作为下一结点的输入，反向传播过程中需要使用正向传播过程中的输出。通常来说，正向传播中计算结点产生的中间特征数据占用的内存空间大，使用间隔时间长，是造成深度神经网络模型训练推理时内存资源紧张的主要原因之一。针对这一问题，业界通常使用数据交换(swap)和重计算(recomputation)等方法降低神经网络训练时的GPU内存消耗。数据交换方法将CPU的DRAM视为临时和外部存储器，把中间特征图交换到CPU内存中，在反向传播过程中再交换回GPU内存中，但是数据交换方法需要额外的通信资源，因此更适合于优化较小的中间数据。重计算方法是指在模型训练过程中特征映射这样的中间结果在正向传播过程中及时地释放，在反向传播的计算需要用到特征映射时，再通过重新计算的方式生成，该方法引入了额外的计算时间，所以更适合优化计算代价较低的中间数据，如Add，Mul等计算结点产生的中间特征数据。

正是因为上述的两种优化策略都存在一些限制，所以最新的方法通常同时使用数据交换和重计算策略对模型训练过程中的GPU内存进行优化。Capuchin方法通过启发式算法比对模型训练中各中间特征数据使用重计算和数据交换的成本，精确的为各数据选取合适的内存优化策略；HOME方法使用粒子群方法为模型中各计算结点分配重计算或数据交换策略以节约GPU内存；Patil等人提出来POET方法以重计算策略和数据交换策略在模型中的应用规范和设备内存为约束条件，使用混合整数线性规划方法求解计算结点策略分配的最优解。但是上述方法通常只考虑对数据选择何种内存优化策略，而不考虑策略执行的起止时间，且这些方法普遍只聚焦于内存资源的优化，并不考虑计算，通信等其他指标，导致模型运行时系统吞吐量低。

发明内容

本发明针对神经网络训练内存优化问题，设计并实现了一种采用强化学习实现AI模型训练内存优化的方法。使用该方法为神经网络训练过程中产生的中间特征数据搜索分配合适的内存优化策略，方法基于神经网络模型结构确定内存优化策略空间，构建用于模拟神经网络训练过程的多维状态矩阵；采用双智能体强化学习的方法协同进行策略搜索；根据策略搜索结果的奖励值，多轮迭代优化以得到收益最大的策略搜索结果，在达到计算加速设备内存优化目标的同时，避免密集化的通信和计算资源消耗；根据策略搜索结果修改神经网络模型的图结构，减少模型训练过程中GPU内存消耗，从而使得单个GPU设备可以处理更大规模的神经网络模型或者更大batch-size的数据样本。

本发明提供一种采用强化学习实现AI模型训练内存优化的方法，步骤如下：

步骤1：收集待优化神经网络模型的基础信息和执行信息，确定内存优化策略空间。

子步骤1.1：使用数据采集工具收集模型的基础信息，包括各个中间特征数据的大小，数据依赖关系，以及模型的执行信息，包括数据结点在实际运行中的GPU访问时间，每个结点GPU内核的计算时间，GPU设备和CPU设备之间的数据传输通信带宽。

子步骤1.2：根据收集的模型信息，过滤输出数据满足内存优化条件的计算结点F＝{F₁,F₂,..F_i}，内存优化策略空间为{A＝{A₁,A₂,..A_i},T_C＝{T_C₁,T_C₂,..,T_C_i}}，即为对第i个结点选择内存优化策略A_i∈{0,1,2}和内存优化策略执行的时间间隔T_C_i∈{c₁,c₂,..,c_k}。

步骤2：构建多维状态矩阵模拟待优化神经网络模型的训练过程，采用双智能体强化学习的方法进行协同策略搜索。

子步骤2.1：初始化两个强化学习智能体action_agent和time_agent，强化学习智能体包含策略选择网络均为包含隐层的前馈神经网络，智能体的输入是当前的多维状态矩阵S_i，输出的是智能体对各个策略选择预测的奖励值P_R_i，选择预测奖励值最高的策略作为智能体为结点F_i选择的策略。

子步骤2.2：分析影响神经网络执行性能的因素，构建多维状态矩阵S(S_mem,S_com,S_comm)，分别用于表示待优化神经网络模型训练过程中的内存资源使用状态，计算资源使用状态和通信资源使用状态。

子步骤2.3：遍历优化结点集合F，使用强化学习智能体agent为结点F_i选择内存优化策略内存优化策略A_i和策略执行时间间隔T_C_i，其中action_agent用于对结点F_i的执行的内存优化策略A_i进行搜索，选择将F_i的计算数据保留在GPU内存中(A_i＝0)，或者使用数据交换(A_i＝1)，重计算(A_i＝2)策略进行内存优化。time_agent用于对结点F_i实际执行优化策略的时间间隔T_C_i进行搜索，在结点F_i的计算数据两次调用间选择最佳时间段c_k(T_C_i＝c_k)执行内存优化策略。

步骤3：动态更新多维状态矩阵，计算策略奖励值，多轮迭代优化强化学习智能体得到收益最大的策略搜索结果。

子步骤3.1：将步骤2中选择结点F_i的内存优化策略A_i和策略执行时间间隔T_C_i，将多维状态矩阵更新为S_i+1，并综合结点信息和状态矩阵表示的内存，计算通信资源等计算策略选择的实际奖励值R_i。

子步骤3.2：，为了分离强化学习的采样和模型优化过程，完全拷贝一份策略选择网络为副本/> 负责进行策略选择以及采集策略执行的奖励值和状态样本。副本/>负责根据策略的奖励值R，不断学习更新策略选择网络的参数，定期同步副本参数/>达到设置轮数，或策略选择的实际奖励值R不再变化时停止迭代，记录迭代过程中实际奖励值R最大的内存优化策略组合{A,T_C}。

步骤4：根据选择的策略修改待优化神经网络模型的结构，减少模型训练过程中的GPU内存消耗。根据步骤3中选择的内存优化策略组合{A,T_C}，对神经网络的图结构进行修改，添加数据换入，换出，重计算结点以减少计算结点的GPU内存消耗，添加trigger结点用于控制内存优化结点的具体执行时间，使用修改后的神经网络模型在真实环境中执行模型训练过程。

本发明具有的有益效果是：

(1)使用多维状态矩阵用于表示神经网络训练中资源变化过程，能模拟多种复杂神经网络模型的训练过程，大大拓展了搜索策略的移植能力。

(2)引入双智能体同时对计算结点的内存优化策略和内存优化策略执行时间间隔进行搜索，在减小GPU内存使用的同时有效提高神经网络训练过程中的系统吞吐量。

(3)通过动态模拟资源使用情况的多维状态矩阵来计算奖励值，能够加快搜索策略的执行和性能评测速度，提高采样效率，从而减少策略整体搜索时间。

(4)在为神经网络添加内存优化结点的同时添加其对应的执行控制依赖,以精确操纵优化结点的具体执行时间，避免资源利用冲突，提高神经网络训练速度。

附图说明

图1是整体架构图；

图2是抽象待优化神经网络模型的计算依赖关系示意图；

图3是构建内存优化策略空间示意图；

图4是多维状态矩阵更新示意图；

图5是待优化神经网络模型的结构修改示意图。

具体实施方式

下面将结合附图和具体实施步骤对本发明做出进一步说明：

如图1至图5所示的一种采用强化学习实现AI模型训练内存优化的方法，包括以下步骤：

步骤1：收集待优化神经网络模型的基础信息和执行信息，确定内存优化策略空间。

子步骤1.1：在真实环境中运行需要进行优化的神经网络，收集神经网络模型的基础信息和执行信息。使用数据交换的方法将神经网络训练过程中的所有中间数据在使用完成后都换出GPU内存，神经网络在训练的过程中GPU内存中只包含当前计算结点所需的数据，这样可以保证当前GPU设备可以对大batch-size的输入数据进行训练而不出现内存不足的错误。使用数据采集工具收集记录模型中各个结点输出的中间特征数据占用的内存空间mem_i,结点间的数据依赖关系；以及训练过程中各个结点的GPU访问时间，计算内核执行时间com_i，GPU设备和CPU设备间的数据传输通信带宽BW。

子步骤1.2：根据收集的待优化神经网络模型信息，使用模型分析矩阵表示模型中的结点数据的计算和依赖关系。设矩阵的大小为2n*2n，n为模型正向计算结点数，横坐标由待优化神经网络中的正向结点F＝{F₁,F₂,..F_n}和反向结点B＝{B_n,..,B₂,B₁}组成，按照结点的逻辑拓扑排序，纵坐标以各结点的计算时间为单位T＝{T₁,T₂,..T_2n}，即在T₁时执行F₁结点的计算任务，T_2n时执行B₁结点的计算任务。图2以一个简单的线性网络结构为例，包含5个正向传播计算结点和5个反向传播计算结点，图中以虚线框标识了结点输出数据被访问的情况。通过分析计算图中的数据访问情况，可以得到输出数据被多次调用的结点，和相邻两次调用的最大访问间隔c_max，通过公式(1)判断结点是否具有优化空间，将符合条件的结点添加到内存优化结点集合F＝{F₁,F₂,..,F_i}。

对集合F中的第i个结点，选择在该两次数据访问时间间隔内执行数据交换(A_i＝1)或者重新计算操作(A_i＝2)，或者驻留在GPU内存中不进行内存优化(A_i＝0)，构成了F_i结点的内存优化策略选择范围A_i∈{0,1,2}。同时也对策略实际执行的起止时刻进行搜索。以图2中的F₁结点为例，结点的最大访问间隔c_max＝[3,8],也就是说在可以在T₃将该数据换出或者丢弃，在T₈将该数据换入或重新计算，然而T₃或T₈可能会存在多个数据交换操作，此时会产生CPU和GPU之间通信拥塞的情况，延长模型训练时间。在该种情况下，我们使用优化间隔为c＝[4,8]或者c＝[3,7]会取得更佳的优化效果，因此对优化时间间隔也进行了变化组合，将组合添加到内存优化策略执行时间间隔选择范围T_C_i中，T_C_i∈{c₁,c₂,..,c_k}。在搜索过程中，为了防止优化时间间隔的组合太多导致搜索空间过大，添加了对于时间间隔的长度限制，限制c_k的长度不小于c_max长度的γ倍，γ∈(0,1]，则时间组合最多为k个。当γ设置为1时，k＝1，时间间隔搜索空间T_C_i＝{c_max}。

步骤2：构建多维状态矩阵模拟待优化神经网络模型训练过程，采用双智能体强化学习的方法进行协同策略搜索。

子步骤2.1：初始化两个强化学习智能体action_agent和time_agent，构建策略选择网络根据输入状态S_i对结点F_i进行策略选择，/>是包含有大小分别为64和128两层隐藏层和一个Flatten层和全联接层的前馈神经网络，其中action_agent包含策略选择网络/>输入是模拟当前训练过程的状态矩阵S_i，输出结果为大小为3的矩阵，分别对应为结点F_i选择不进行内存优化、数据交换、重计算这三种策略的预测奖励值P_R_i。而time_agent包含的/>输入是模拟当前训练过程的状态矩阵S_i，输出结果为大小为k的矩阵，代表为优化策略选择不同执行时间间隔的预测奖励值P_R_i。

子步骤2.2：根据步骤1中收集的神经网络信息，构建多维状态矩阵S(S_mem,S_com,S_comm)用于模拟模型的训练过程中的计算、内存和通信资源使用状态。状态矩阵的大小和坐标轴同步骤1中的模型分析矩阵一致，在执行策略选择前初始化状态矩阵，计算状态矩阵S_com中(F_i/B_2n-i,T_i)位置保存对应结点的GPU计算执行时间。如图3和图4所示，如果神经网络中一个结点计算输出的数据会被多次调用，那么从数据产生到最后一次调用结束，数据所占用的计算内存都是固定的。假设在一轮训练过程中，结点F₁的输出最后被B₁结点调用，那么在内存状态矩阵S_mem中，从(F₁,T₂)到(F₁,T_2n)的值都是F₁结点的输出数据占用的GPU内存大小。通信状态矩阵S_comm用于表示结点的通信状态，当结点F_i由于数据交换策略产生GPU-CPU的数据通信时，在矩阵中记录通信状态变化，值为交换的数据内存大小mem_i除以GPU和CPU之间的数据传输通信带宽BW，即数据通信时间。

子步骤2.3：遍历优化结点集合F，action_agent用于对结点F_i的执行的内存优化策略A_i进行搜索，根据策略选择网络预测的奖励值P_R_i，选择奖励值值最大的策略作为结点F_i的内存优化策略A_i，根据内存优化策略A_i的值选择将F_i的计算数据保留在GPU内存中(A_i＝0)，或者使用数据交换(A_i＝1)，重计算(A_i＝2)策略进行内存优化。time_agent用于对结点F_i实际执行优化策略的时间间隔T_C_i进行搜索，在结点F_i的计算数据两次调用间选择最佳时间段c_k执行内存优化策略。每遍历一轮优化结点集合F，即完成一次待优化神经网络模型的内存优化策略分配。

步骤3：动态更新多维状态矩阵，计算策略奖励值，多轮迭代优化强化学习智能体得到收益最大的策略搜索结果。

子步骤3.1：在实际训练过程中，将结点从GPU中卸载再重新生成会改变训练过程中的资源使用情况。同理，在强化学习智能体action_agent和time_agent在完成对结点的策略选择之后，用于模拟神经网络训练过程的多维状态矩阵也会对资源的使用情况做同步的更新，更新规则如下：(1)策略选择为GPU(A_i＝0)时，S_i+1＝S_i。(2)策略选择为数据交换(A_i＝1)时，则需要对数据先换出到CPU内存，再换回GPU内存。在通信状态矩阵S_comm上，T_C_i的起始时(T_C_i.start)和结束时(T_C_i.end)添加F_i的数据通信代价，而在t∈T_C_i，因为数据已经换出到CPU内存中，则将内存状态矩阵S_mem中对应时间段的内存消耗设置为0。(3)策略选择为重计算(A_i＝2)时，会在调用F_i的输出数据之前对F_i进行重新计算，在计算状态矩阵S_com上，在T_C_i的结束时(T_C_i.end)添加F_i的重新计算代价，而在t∈T_C_i，因为数据已经丢弃，将内存状态矩阵S_mem中对应时间段的内存消耗设置为0。

多维状态矩阵更新的细节过程如下：

计算F_i策略选择的实际奖励值R_i，奖励值对强化学习智能体进行策略搜索具有指导作用，对奖励值的计算也要综合考虑策略在实际训练过程中对多种资源的影响。R_i的计算遵循以下规则：

(1)内存资源判断。对当前的内存状态矩阵S_mem，按纵坐标计算训练内存峰值，保留纵坐标得到各个结点MEM_T＝{MEM_T₁,MEM_T₂,..MEM_T_2n}，计算T_i的正向结点F和反向结点B所占用的内存之和，通常来说，在反向传播刚刚开始时待优化神经网络模型的内存占用最大，此时正向传播计算结点的内存占用未被释放，因此为内存消耗的峰值。如果内存消耗的峰值max(MEM_T)小于设备内存消耗，说明GPU设备内存可以支撑待优化神经网络模型的训练过程，则中断该轮搜索，并且对该轮策略搜索的所有优化策略选择增加奖励值。如果内存峰值消耗大于设备内存，说明仍需要从GPU内存中卸载部分数据。此时如果A_i＝1or2，并且T_C_i包含内存消耗的峰值时刻，则该策略选择降低了待优化神经网络模型训练时峰值内存消耗，所以会为策略选择添加奖励值。

MEM_T_i＝Σ_1≤j≤n S_mem[F_j/B_j.node_id][i] (3)

(2)计算资源判断。对计算资源的判断主要用于判断重计算策略选择的时间间隔是否合理，重计算策略需要在被反向调用之前对结点进行重新计算，执行重新计算的时刻即为T_C_i.end，比较F_i的计算时间和的计算时间，如果重F_i的计算时间较长，则该策略选择会延长待优化神经网络模型的训练时间，因此减少对应策略的奖励值。

(3)通信资源判断。对通信资源的判断主要用于判断数据交换策略选择的时间间隔是否合理，数据交换策略包含数据换出和数据换入两个动作，分别在T_C_i.start将数据换出，在T_C_i.end将数据换入，如果两个时刻其他计算结点存在通信资源消耗则减少对应策略的奖励值。

(4)优化时间间隔判断。不论是使用数据交换方法或者重新计算方法进行内存优化时，都需要额外的通信或是计算，对于优化时间间隔较小的结点来说，可能数据刚刚换出计算加速设备内存，又需要重新载入，因此需要将这种结点输出数据尽量保存在GPU内存，所以当A_i＝0时，优化时间间隔越短，增加奖励值越高。而当A_i＝1或者2时，减少奖励值。

(5)计算执行依赖判断。当A_i＝2，这要求重计算的结点的输入也在内存中，否则无法进行重新计算。所以F_i选择了重新计算策略，那么会校验对结点的输入是否在内存中，如果在内存中，增加策略奖励，如果不在内存中，则会减少该策略的奖励值，同时也会减少在经验回放池中的输入结点的策略奖励值。

子步骤3.2：为了分离强化学习的采样和策略选择网络优化过程，完全拷贝一份策略选择网络为副本/> 负责进行策略选择以及采集策略执行的奖励值和状态样本。副本/>负责使用经验回放方法对参数进行学习和更新。将(S_i,A_i,T_C_i,R_i,S_i+1)保存为一条transition，构建经验回放池保存N条transition:(S_i,A_i,T_C_i,R_i,S_i+1)，(S_i+1,A_i+1,T_C_i+1,R_i+1,S_i+2)，...，(S_i+N,A_i+N,T_C_i+N,R_i+N,S_i+N+1)，每隔一段时间从N条transition中抽取batch_size条数据对/>的参数进行更新，通过将标签值和/>预测的奖励值P_R_i进行比较，使用公式5计算出loss值，然后根据loss值使用梯度下降的方法更新副本/>网络的参数值，定期同步副本参数

loss_i＝(R_i+γ*P_R_i)-P_R_i (5)

当迭代才开始，用于策略选择的的策略选择具有相当大的随机性，然而随着不断更新同步网络参数，/>的选择结果会不断趋于实际奖励值更大的策略，从而得到结果最优的策略组合。当达到设置的遍历轮数round_num，或连续10轮策略选择的实际奖励值R不再变化时停止迭代，返回实际奖励值最大的策略组合{A,T_C}。其中，round_num的数值根据经验设定。

步骤2，步骤3中介绍的强化学习策略搜索算法的主体过程如下：

步骤4：根据选择的策略修改待优化神经网络模型的图结构，减少模型训练过程中的GPU内存消耗。

遍历结点集合F，根据步骤3中最优策略组合{A,T_C}对神经网络模型图进行修改，修改规则如下：

(1)A_i＝0时，不对计算图中的结点F_i作修改。

(2)A_i＝1时，对结点进行数据交换图修改，首先需要断开F_i结点输出数据和反向传播中对应结点的数据依赖，然后新增内存优化结点：swap_out和swap_in，swap_out结点负责将数据从GPU内存转移到CPU内存中，swap_in结点负责将数据从CPU内存中转移回GPU内存，在swap_out结点执行完成到swap_in结点执行完成的这一段时间内，计算结点F_i的输出数据不再占用GPU内存。添加执行控制依赖：swap_out_trigger和swap_in_trigger，swap_out_trigger结点在T_C_i.start开始执行，swap_in_trigger结点在T_C_i.end开始执行。如图5中对应虚线所示，swap_out结点的输入是F_i结点的输出特征数据和swap_out_trigger结点，swap_in结点的输入是swap_out结点和swap_in_trigger结点，swap_in结点的输出则是反向传播中的调用F_i的输出数据的计算结点。

(3)A_i＝2时，对结点进行重新计算的图修改，首先断开F_i结点输出数据和反向传播中调用结点的数据依赖关系，在断开断开F_i结点输出数据和反向传播中调用结点的数据依赖关系之后，F_i结点输出数据数据会自动被内存管理系统回收，不再占用GPU内存空间。然后新增内存优化结点：recomputation，recomputation结点负责重新执行F_i结点的计算过程并生成输出数据。添加执行控制依赖：recomputation_trigger，recomputation_trigger结点在T_C_i.end开始执行。如图5中对应虚线所示，recomputation结点的输入是原结点的输入和recomputation_trigger结点，recomputation结点的输出则为反向传播中调用F_i的输出数据的计算结点。

最后，使用修改后的神经网络模型图在真实环境中进行训练。

Claims (6)

1.一种采用强化学习实现AI模型训练内存优化的方法，其特征在于包括以下步骤：

步骤1，收集待优化神经网络模型的基础信息和执行信息，使用模型分析矩阵表示待优化神经网络模型中计算结点的数据依赖关系，筛选满足内存优化条件的计算结点，确定内存优化策略空间，策略空间包含内存优化策略选择范围和内存优化策略执行时间间隔选择范围；

步骤2，构建多维状态矩阵模拟待优化神经网络模型的训练过程，基于模拟的多维状态矩阵，采用双智能体对待优化神经网络模型中各个计算结点进行搜索；

搜索包括：内存优化策略和内存优化策略执行时间间隔；

步骤3，根据选择的策略更新状态矩阵，动态模拟实际训练环境中的资源变化；根据策略对训练资源的影响计算奖励值，并根据策略奖励值调整智能体参数，通过多轮迭代优化强化学习智能体得到奖励值最大的策略搜索结果；

步骤4，根据最优策略搜索结果修改待优化神经网络模型的图结构，为待优化神经网络的计算结点添加强化学习方法所选择的内存优化策略；

在真实训练环境中执行修改后的计算图，减少神经网络模型在训练过程中的GPU内存消耗。

2.根据权利要求1中所述的一种采用强化学习实现AI模型训练内存优化的方法，其特征在于，同时对结点执行的内存优化策略和策略执行时间间隔进行搜索；步骤1中满足内存优化条件的计算结点包括：

计算结点的输出数据需被多次调用，且相邻两次调用需满足条件：调用时间间隔需大于对该结点使用内存优化策略的执行时间，使用以下公式表达

其中，mem_i代表结点的输出数据内存占用，com_i代表结点的计算时间，c_max代表结点输出数据相邻两次调用的最大时间间隔，BW为GPU和CPU间的通信带宽。

3.根据权利要求1中所述的一种采用强化学习实现AI模型训练内存优化的方法，其特征在于，步骤2中构建多维状态矩阵模拟待优化神经网络模型的训练过程是指：通过矩阵模拟真实环境中的资源使用情况，用大小为3*2n*2n的矩阵表示模型训练过程中各计算结点的内存占用，计算时间和通信时间，其中，n为被优化神经网络模型中的正向计算结点数。

4.根据权利要求1中所述的一种采用强化学习实现AI模型训练内存优化的方法，其特征在于，步骤3中更新状态矩阵是指，根据计算结点选择的内存优化策略更新状态矩阵中对应数据值；

内存优化策略为数据交换时，更新矩阵中计算结点的内存占用和通信时间；

内存优化策略为重计算时，更新矩阵中计算结点的内存占用和计算时间；

通过更新状态矩阵取代重新收集数据的过程，从而加快策略搜索迭代速度。

5.根据权利要求1中所述的一种采用强化学习实现AI模型训练内存优化的方法，其特征在于，步骤3中根据策略对训练资源的影响计算奖励值是指，通过内存资源判断、计算资源判断、通信资源判断、优化时间间隔判断、计算执行依赖判断这五种规则来计算策略奖励值。

6.根据权利要求1中所述的一种采用强化学习实现AI模型训练内存优化的方法，其特征在于，步骤4中的计算图修改方法，根据最优策略选择添加对应的内存优化结点，包括swap_in结点、swap_out结点和recomputation结点，同时添加其对应的执行控制依赖用于控制内存优化结点执行的起止时间，包括swap_in_trigger、swap_in_trigger和recomputation_trigger，在真实训练环境中执行修改后的计算图。