numpy - 将 numpy 数组运算符保留在 np.float32 中

Question

a = np.random.random((512,512,3)).astype(np.float32)
b = np.ones((512,512,1), dtype=np.int32)
c = a / b
c.dtype

>> dtype('float64')

将一个 float32 矩阵除以一个 int32 矩阵得到一个 float64 矩阵。目前我必须做的

return c.astype(np.float32)

这是 CPU 的额外工作。有没有办法避免最终转换并告诉 numpy 在 float32 中完成工作？

Answer 1

您必须使用 np.divide() 的 out 参数。

[~/scratch]
|1> a = np.random.random((512,512,3)).astype(np.float32)

[~/scratch]
|2> b = np.ones((512,512,1), dtype=np.int32)

[~/scratch]
|3> c = np.empty_like(a)

[~/scratch]
|4> c = np.divide(a, b, c)

[~/scratch]
|5> c.dtype
dtype('float32')

在 numpy 1.6 中，您将能够执行 c = np.divide(a, b, dtype=np.float32)