首页 > 解决方案 > NumPy 中有 unstack 吗?

问题描述

NumPy中有np.stack,但是否有与tf.unstack相同的对立面?np.unstack

标签: pythonnumpy

解决方案


遇到这么晚,这里有一个更简单的答案:

def unstack(a, axis=0):
    return np.moveaxis(a, axis, 0)
#    return list(np.moveaxis(a, axis, 0))

作为奖励,结果仍然是一个 numpy 数组。如果您只是 python-unwrap 它会发生解包:

A, B, = unstack([[1, 2], [3, 4]], axis=1)
assert list(A) == [1, 3]
assert list(B) == [2, 4]

不出所料,它也是最快的:

❯ python -m timeit -s "import numpy as np; a=np.array(np.meshgrid(np.arange(1000), np.arange(1000)));" "C = [np.squeeze(e, 1) for e in np.split(a, a.shape[1], axis = 1)]"
100 loops, best of 5: 2.64 msec per loop
    
❯ python -m timeit -s "import numpy as np; a=np.array(np.meshgrid(np.arange(1000), np.arange(1000)));" "C = [np.take(a, i, axis = 1) for i in range(a.shape[1])]"       
50 loops, best of 5: 5.08 msec per loop

❯ python -m timeit -s "import numpy as np; a=np.array(np.meshgrid(np.arange(1000), np.arange(1000)));" "C = np.moveaxis(a, 1, 0)"
100000 loops, best of 5: 3.89 usec per loop

# With list conversion
❯ python -m timeit -s "import numpy as np; a=np.array(np.meshgrid(np.arange(1000), np.arange(1000)));" "C = list(np.moveaxis(a, 1, 0))"
1000 loops, best of 5: 205 usec per loop

推荐阅读