🐍

共有メモリに numpy.ndarray を載せる

2021/02/03に公開

TL; DR

Python 3.7 対応
機械学習で利用する16bit 浮動小数点 (np.float16)対応

from multiprocessing.sharedctypes import RawArray
import numpy as np

# 載せたい型とサイズの指定
dtype = np.float16
shape = (3,2)

# ここから
try:
    ctype = np.ctypeslib.as_ctypes_type(np.dtype(dtype))
except NotImplementedError:
    for d in (np.int8, np.int16, np.int32, np.int64):
        _d = np.dtype(d)
	if np.dtype(dtype).itemsize == _d.itemsize:
	    ctype = np.ctypeslib.as_ctypes_type(_d)
	    break
    else:
        raise

len = int(np.array(shape, copy=False,dtype="int").prod())
data = np.ctypeslib.as_array(RawArray(ctype, len))
data.shape = shape
array = data.view(dtype)

0. モチベーション

個人開発している強化学習におけるExperience Replayライブラリで、マルチプロセス対応をするために、 numpy.ndarray のデータを共有メモリに配置したかった。

1. 共有メモリの作り方

Python 3.8 以降はバイト数を直接指定できる multiprocessing.shared_memory が標準ライブラリに追加されているが、 Python 3.7 では利用できない。

Python 3.7 では multiprocessing.sharedctypes.RawArray を利用することができる。
これは、 ctypes に定義されている「C互換の型または 1文字の型コード」と「要素数」で指定する必要がある

2. np.ctypeslib の利用

numpy.dtype から ctypes に変換するのは、 numpy.ctypeslib の関数群を利用することができる。

これで完成したと思っていた。

from multiprocessing.sharedctypes import RawArray
import numpy as np

# 載せたい型とサイズの指定
dtype = np.single
shape = (3,2)

# ここから
ctype = np.ctypeslib.as_ctypes_type(np.dtype(dtype))
len = int(np.array(shape, copy=False,dtype="int").prod())
data = np.ctypeslib.as_array(RawArray(ctype, len))
data.shape = shape # reshape を使うとコピーしうるので、attribute に直代入

3. C非互換型への対応

完成したと思っていたらユーザーから isseue 報告がやってきた。「np.float16」でエラーになると。

NotImplementedError: Converting dtype('float16') to a ctypes type
16bit浮動小数点は、昨今の64bit環境ではC言語に互換型が無いので、変換できずにエラーとなっていた。

機械学習では精度を犠牲にしてデータサイズを小さくすることも多々あるので、 np.float16 非対応だとまずいと思い対策を検討した。

ctypes への変換を経ずに、同じバイト数の共有メモリを無理やり読み替える方法を採用することにした。 ndarray の reintepret は view() に dtype を指定することで可能である。

ctypes に変換できない時に、変数のサイズが 8bit、16bit、32bit、64bit であるかをチェックして該当する型で共有メモリを作成する。 (128bit は環境依存なので対応させていない。)

try:
    ctype = np.ctypeslib.as_ctypes_type(np.dtype(dtype))
except NotImplementedError:
    for d in (np.int8, np.int16, np.int32, np.int64):
        _d = np.dtype(d)
	if np.dtype(dtype).itemsize == _d.itemsize:
            ctype = np.ctypeslib.as_ctypes_type(_d)
            break
    else:
        # どれにもマッチしなければ、再度 NotImplementedError を上げる
        raise

いっそのこと全部 np.byte の倍数でメモリを確保しても良いが、 np.byte もサイズが環境依存で 1byte とも限らないので、可能な限り実装済みの as_ctypes_type にまかせて拾えない部分だけ個別対応することとした。

完成版 (再掲)

from multiprocessing.sharedctypes import RawArray
import numpy as np

# 載せたい型とサイズの指定
dtype = np.float16
shape = (3,2)

# ここから
try:
    ctype = np.ctypeslib.as_ctypes_type(np.dtype(dtype))
except NotImplementedError:
    for d in (np.int8, np.int16, np.int32, np.int64):
        _d = np.dtype(d)
	if np.dtype(dtype).itemsize == _d.itemsize:
	    ctype = np.ctypeslib.as_ctypes_type(_d)
	    break
    else:
        raise

len = int(np.array(shape, copy=False,dtype="int").prod())
data = np.ctypeslib.as_array(RawArray(ctype, len))
data.shape = shape
array = data.view(dtype)