如何在Python中用OpenMP或类似技术并行化for循环？

问题背景 #

需要并行化Replay Buffer中_encode_sample方法里的for循环：该循环迭代idxes（共1024次），每次从self._storage中取出数据并整理成最终数组，希望借助OpenMP或类似并行技术加速执行。

原始串行代码 #

def _encode_sample(self, idxes):
    obses_t, actions, rewards, obses_tp1, dones = [], [], [], [], []
    for i in idxes:
        data = self._storage[i]
        obs_t, action, reward, obs_tp1, done = data
        obses_t.append(np.array(obs_t, copy=False))
        actions.append(np.array(action, copy=False))
        rewards.append(reward)
        obses_tp1.append(np.array(obs_tp1, copy=False))
        dones.append(done)
    return np.array(obses_t), np.array(actions), np.array(rewards), np.array(obses_tp1), np.array(dones)

已有的Ray优化方案 #

以下是基于Ray实现的并行版本：

import numpy as np
import random
import ray
import psutil

num_cpus = psutil.cpu_count(logical=False)
ray.init(num_cpus=1)

class ReplayBuffer(object):
    def __init__(self, size):
        """Create Prioritized Replay buffer.

        Parameters
        ----------
        size: int
            Max number of transitions to store in the buffer. When the buffer
            overflows the old memories are dropped.
        """
        self._storage = []
        self._maxsize = int(size)
        self._next_idx = 0

    def __len__(self):
        return len(self._storage)

    def clear(self):
        self._storage = []
        self._next_idx = 0

    def add(self, obs_t, action, reward, obs_tp1, done):
        data = (obs_t, action, reward, obs_tp1, done)
        if self._next_idx >= len(self._storage):
            self._storage.append(data)
        else:
            self._storage[self._next_idx] = data
        self._next_idx = (self._next_idx + 1) % self._maxsize

    def _encode_sample(self, idxes):
        n = 256
        split_idxes = [idxes[i * n:(i + 1) * n] for i in range((len(idxes) + n - 1) // n )]
        futures = []
        # 先将存储对象放入Ray对象存储，避免重复序列化
        storage_ref = ray.put(self._storage)
        for subrange in split_idxes:
            futures.append(_encode_sample_helper.remote(storage_ref, subrange))

        obses_t, actions, rewards, obses_tp1, dones = [], [], [], [], []
        outputs = ray.get(futures)

        for a, b, c, d, e in outputs:
            obses_t.extend(a)
            actions.extend(b)
            rewards.extend(c)
            obses_tp1.extend(d)
            dones.extend(e)

        return np.array(obses_t), np.array(actions), np.array(rewards), np.array(obses_tp1), np.array(dones)

    def make_index(self, batch_size):
        return [random.randint(0, len(self._storage) - 1) for _ in range(batch_size)]

    def make_latest_index(self, batch_size):
        idx = [(self._next_idx - 1 - i) % self._maxsize for i in range(batch_size)]
        np.random.shuffle(idx)
        return idx

    def sample_index(self, idxes):
        return self._encode_sample(idxes)

    def sample(self, batch_size):
        """Sample a batch of experiences.

        Parameters
        ----------
        batch_size: int
            How many transitions to sample.

        Returns
        -------
        obs_batch: np.array
            batch of observations
        act_batch: np.array
            batch of actions executed given obs_batch
        rew_batch: np.array
            rewards received as results of executing act_batch
        next_obs_batch: np.array
            next set of observations seen after executing act_batch
        done_mask: np.array
            done_mask[i] = 1 if executing act_batch[i] resulted in
            the end of an episode and 0 otherwise.
        """
        if batch_size > 0:
            idxes = self.make_index(batch_size)
        else:
            idxes = range(0, len(self._storage))
        return self._encode_sample(idxes)

    def collect(self):
        return self.sample(-1)

@ray.remote(num_cpus=1)
def _encode_sample_helper(_storage, subrange):
    obses_t, actions, rewards, obses_tp1, dones = [], [], [], [], []
    for i in subrange:
        data = _storage[i]
        obs_t, action, reward, obs_tp1, done = data
        obses_t.append(np.array(obs_t, copy=False))
        actions.append(np.array(action, copy=False))
        rewards.append(reward)
        obses_tp1.append(np.array(obs_tp1, copy=False))
        dones.append(done)
    return np.array(obses_t), np.array(actions), np.array(rewards), np.array(obses_tp1), np.array(dones)

基于OpenMP风格的替代方案 #

Python原生不支持OpenMP，但可以通过Numba实现类似OpenMP的轻量级线程并行（Numba会绕过GIL，底层使用OpenMP作为并行后端）：

Numba并行实现示例 #

import numba as nb
import numpy as np

class ReplayBuffer(object):
    # ... 其他方法保持不变 ...

    def _encode_sample(self, idxes):
        # 将storage转换为NumPy数组（假设存储的元素结构一致）
        storage_np = np.array(self._storage, dtype=object)
        # 调用Numba并行处理函数
        return _encode_sample_numba(storage_np, np.array(idxes))

@nb.njit(parallel=True)
def _encode_sample_numba(storage, idxes):
    n = len(idxes)
    # 假设obs_t是固定形状的数组，先获取形状初始化结果数组
    obs_shape = storage[idxes[0]][0].shape
    action_shape = storage[idxes[0]][1].shape
    
    obses_t = np.empty((n,) + obs_shape, dtype=storage[idxes[0]][0].dtype)
    actions = np.empty((n,) + action_shape, dtype=storage[idxes[0]][1].dtype)
    rewards = np.empty(n, dtype=np.float32)
    obses_tp1 = np.empty((n,) + obs_shape, dtype=storage[idxes[0]][3].dtype)
    dones = np.empty(n, dtype=np.bool_)
    
    # 使用prange实现并行循环
    for i in nb.prange(n):
        data = storage[idxes[i]]
        obses_t[i] = data[0]
        actions[i] = data[1]
        rewards[i] = data[2]
        obses_tp1[i] = data[3]
        dones[i] = data[4]
    
    return obses_t, actions, rewards, obses_tp1, dones

关键说明 #

• Numba的@nb.njit(parallel=True)会自动启用OpenMP并行，prange标记并行循环。
• 提前初始化结果数组而非动态append，能进一步提升效率，避免Python列表的扩容开销。
• 若self._storage中的元素是异构结构，可能需要调整数组初始化方式。

方案选择建议 #

• 若追求低开销的线程级并行（类似OpenMP），优先选Numba实现，代码改动小，性能提升明显。
• 若需要分布式并行或跨节点调度，Ray方案更合适，但注意原代码中ray.init(num_cpus=1)限制了CPU使用，应改为ray.init(num_cpus=psutil.cpu_count(logical=True))以利用全部核心。
• 轻量级多进程并行也可使用multiprocessing.Pool，但序列化开销略高于Ray的对象存储机制。

内容的提问来源于stack exchange，提问作者Sherlock