深度强化学习训练优化：如何平衡CPU和GPU使用

创作时间:

作者:

@小白创作中心

深度强化学习训练优化：如何平衡CPU和GPU使用

引用

CSDN

https://blog.csdn.net/weixin_56760882/article/details/139098003

在深度强化学习训练中，如何平衡CPU和GPU的使用以达到最佳训练速度是一个值得探讨的问题。本文通过实验和分析，提出了CPU采样+GPU训练模型的方法，并提供了具体的实现逻辑和代码示例。

问题背景

在深度强化学习中，模型训练通常需要大量的计算资源。传统的做法是使用GPU进行全量训练，但作者在实践中发现，将模型从CPU拷贝到GPU会消耗大量时间，导致整体训练速度变慢。因此，作者开始研究如何更有效地利用CPU和GPU的特性来优化训练过程。

实验环境

算法：PPO（Proximal Policy Optimization）
环境：CartPole-v0
模型：128单隐层
样本数：0（on-line）

原始GPU训练结果

耗时：2分钟8秒
GPU占用：45%左右
CPU占用：27%

CPU训练结果

耗时：47.4秒（经常CPU满载后，CPU会变慢，后续测时为53秒）
CPU占用：满载100%

解决方案

参考文献解释了CPU训练速度快、GPU训练速度慢的原因：模型从CPU拷贝到GPU花了大量时间，且环境交互时的少量运算，CPU比GPU运算速度更快。基于此，作者提出了CPU采样+GPU训练模型的方法。

CPU采样+GPU训练的优势

环境交互方面：强化学习时要进行大量的环境交互，也就是进行文件的读取操作或者少量计算操作。CPU对于文件读取和少量计算往往比GPU更快，更准确，是因为CPU有少量且强大的核心，设计时是为了专注于处理不同的任务。
模型训练方面：模型中有大量的参数，但模型更新时都是简单的矩阵运算。GPU有大量的小核心，专注于图形处理和矩阵运算，在这方面速度比CPU要快。

实现逻辑

模型都在CPU上生成，由CPU与环境交互，由GPU训练模型。

在环境采样前将模型参数上传到（copy到）CPU上
在模型训练前（更新前）将模型copy到GPU上

PPO算法示例代码

from tqdm import tqdm

def train_on_policy_agent(env, agent, num_episodes):
    return_list = []
    for i in range(10):
        with tqdm(total=int(num_episodes/10), desc='Iteration %d' % i) as pbar:
            for i_episode in range(int(num_episodes/10)):
                episode_return = 0
                transition_dict = {'states': [], 'actions': [], 'next_states': [], 'rewards': [], 'dones': []}
                state = env.reset(seed =0)[0] #1.改 gym 0.26.0版本后，env.reset()返回的是一个字典，所以需要加上[0]
                agent.actor.to('cpu')
                done = False
                while not done:
                    action = agent.take_action(state) 
                    next_state, reward,terminated, truncated, _ = env.step(action) #2.改看gym版本0.26.2版本的 
                    done = terminated or truncated
                    transition_dict['states'].append(state)
                    transition_dict['actions'].append(action)
                    transition_dict['next_states'].append(next_state)
                    transition_dict['rewards'].append(reward)
                    transition_dict['dones'].append(done)
                    state = next_state
                    episode_return += reward
                return_list.append(episode_return)
                agent.actor.to('cuda') #
                agent.critic.to('cuda')
                agent.update(transition_dict)
                if (i_episode+1) % 10 == 0:
                    pbar.set_postfix({'episode': '%d' % (num_episodes/10 * i + i_episode+1), 'return': '%.3f' % np.mean(return_list[-10:])})
                pbar.update(1)
    return return_list