在 redis 中存储 numpy 数组的最快方法

Question

我在一个 AI 项目中使用 redis。

这个想法是让多个环境模拟器在大量 cpu 内核上运行策略。模拟器将体验（状态/动作/奖励元组列表）写入 redis 服务器（重播缓冲区）。然后训练过程将体验作为数据集读取以生成新策略。新策略被部署到模拟器，之前运行的数据被删除，流程继续。

大部分体验都是在“状态”中捕获的。这通常表示为一个大型的 numpy 数组，例如 80 x 80。模拟器会以 CPU 允许的速度生成这些数组。

为此，有没有人有好的想法或经验最好/最快/最简单的方法来编写大量的numpy数组到redis。这一切都在同一台机器上，但后来可能在一组云服务器上。欢迎使用代码示例！

Answer 1

我不知道它是否最快，但你可以试试这样的......

将 Numpy 数组存储到 Redis 如下所示 - 参见函数 toRedis()：

• 获取 Numpy 数组的形状并编码
• 将 Numpy 数组作为字节附加到形状中
• 将编码数组存储在提供的密钥下

检索 Numpy 数组的过程如下 - 参见函数 fromRedis()：

• 从 Redis 中检索与提供的键对应的编码字符串
• 从字符串中提取 Numpy 数组的形状
• 提取数据并重新填充 Numpy 数组，重塑为原始形状

---

#!/usr/bin/env python3

import struct
import redis
import numpy as np

def toRedis(r,a,n):
   """Store given Numpy array 'a' in Redis under key 'n'"""
   h, w = a.shape
   shape = struct.pack('>II',h,w)
   encoded = shape + a.tobytes()

   # Store encoded data in Redis
   r.set(n,encoded)
   return

def fromRedis(r,n):
   """Retrieve Numpy array from Redis key 'n'"""
   encoded = r.get(n)
   h, w = struct.unpack('>II',encoded[:8])
   # Add slicing here, or else the array would differ from the original
   a = np.frombuffer(encoded[8:]).reshape(h,w)
   return a

# Create 80x80 numpy array to store
a0 = np.arange(6400,dtype=np.uint16).reshape(80,80) 

# Redis connection
r = redis.Redis(host='localhost', port=6379, db=0)

# Store array a0 in Redis under name 'a0array'
toRedis(r,a0,'a0array')

# Retrieve from Redis
a1 = fromRedis(r,'a0array')

np.testing.assert_array_equal(a0,a1)

您可以通过将 Numpy 数组的 dtype 与形状一起编码来增加灵活性。我没有这样做是因为您可能已经知道您的所有数组都是一种特定类型，然后代码会无缘无故变得更大更难阅读。

现代 iMac 的粗略基准：

80x80 Numpy array of np.uint16   => 58 microseconds to write
200x200 Numpy array of np.uint16 => 88 microseconds to write

---

关键字：Python、Numpy、Redis、数组、序列化、序列化、键、增量、唯一

Answer 2

tobytes() 函数的存储效率不是很高。 为了减少必须写入 redis 服务器的存储空间，您可以使用 base64 包：

def encode_vector(ar):
    return base64.encodestring(ar.tobytes()).decode('ascii')

def decode_vector(ar):
    return np.fromstring(base64.decodestring(bytes(ar.decode('ascii'), 'ascii')), dtype='uint16')

@EDIT：好的，由于 Redis 将值存储为字节字符串，因此直接存储字节字符串的存储效率更高。但是，如果将其转换为字符串、将其打印到控制台或将其存储在文本文件中，则进行编码是有意义的。

Answer 3

您可以查看 Mark Setchell 的回答，了解如何将字节实际写入 Redis。下面我重写了函数 fromRedis 和 toRedis 以考虑可变维度大小的数组并包括数组形状。

def toRedis(arr: np.array) -> str:
    arr_dtype = bytearray(str(arr.dtype), 'utf-8')
    arr_shape = bytearray(','.join([str(a) for a in arr.shape]), 'utf-8')
    sep = bytearray('|', 'utf-8')
    arr_bytes = arr.ravel().tobytes()
    to_return = arr_dtype + sep + arr_shape + sep + arr_bytes
    return to_return

def fromRedis(serialized_arr: str) -> np.array:
    sep = '|'.encode('utf-8')
    i_0 = serialized_arr.find(sep)
    i_1 = serialized_arr.find(sep, i_0 + 1)
    arr_dtype = serialized_arr[:i_0].decode('utf-8')
    arr_shape = tuple([int(a) for a in serialized_arr[i_0 + 1:i_1].decode('utf-8').split(',')])
    arr_str = serialized_arr[i_1 + 1:]
    arr = np.frombuffer(arr_str, dtype = arr_dtype).reshape(arr_shape)
    return arr

Answer 4

您还可以考虑使用msgpack-numpy，它提供“编码和解码例程，可以使用高效的 msgpack 格式对 numpy 提供的数值和数组数据类型进行序列化和反序列化。” -- 见https://msgpack.org/。

快速概念验证：

import msgpack
import msgpack_numpy as m
import numpy as np
m.patch()               # Important line to monkey-patch for numpy support!

from redis import Redis

r = Redis('127.0.0.1')

# Create an array, then use msgpack to serialize it 
d_orig = np.array([1,2,3,4])
d_orig_packed = m.packb(d_orig)

# Set the data in redis
r.set('d', d_orig_packed)

# Retrieve and unpack the data
d_out = m.unpackb(r.get('d'))

# Check they match
assert np.alltrue(d_orig == d_out)
assert d_orig.dtype == d_out.dtype

在我的机器上，msgpack 的运行速度比使用 struct 快得多：

In: %timeit struct.pack('4096L', *np.arange(0, 4096))
1000 loops, best of 3: 443 µs per loop

In: %timeit m.packb(np.arange(0, 4096))
The slowest run took 7.74 times longer than the fastest. This could mean that an intermediate result is being cached.
10000 loops, best of 3: 32.6 µs per loop

Answer 5

试试 Plasma，因为它避免了序列化/反序列化开销。

使用 pip install pyarrow 安装 Plasma

文档：https://arrow.apache.org/docs/python/plasma.html

首先，启动具有 1 GB 内存的 Plasma[终端]：

plasma_store -m 1000000000 -s /tmp/plasma

import pyarrow.plasma as pa
import numpy as np
client = pa.connect("/tmp/plasma")
temp = np.random.rand(80,80)

写入时间：130 µs vs 782 µs（Redis 实现：Mark Setchell 的回答）

使用等离子大页面可以提高写入时间，但仅适用于 Linux 机器：https://arrow.apache.org/docs/python/plasma.html#using-plasma-with-huge-pages

获取时间：31.2 µs vs 99.5 µs（Redis 实现：Mark Setchell 的回答）

PS：代码是在 MacPro 上运行的