python - 如何在 numpy 中沿小维度优化点积?
问题描述
我有np.ndarray
两个
a
是一个形状(13000, 8, 315000)
和类型的数组uint8
b
是一个形状(8,)
和类型的数组float32
我想将沿第二维(8)的每个切片乘以相应的元素b
并沿该维求和(即沿第二轴的点积)。结果将是形状(13000, 315000)
我设计了两种方法来做到这一点:
np.einsum('ijk,j->ik', a, b)
:使用%timeit
它给49 s ± 12.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
np.dot(a.transpose(0, 2, 1), b)
:使用%timeit
它给1min 8s ± 3.54 s per loop (mean ± std. dev. of 7 runs, 1 loop each)
有更快的替代品吗?
补充资料
np.show_config()
返回:
blas_mkl_info:
NOT AVAILABLE
openblas_lapack_info:
libraries = ['openblas', 'openblas']
language = c
library_dirs = ['/usr/local/lib']
define_macros = [('HAVE_CBLAS', None)]
lapack_mkl_info:
NOT AVAILABLE
openblas_info:
libraries = ['openblas', 'openblas']
language = c
library_dirs = ['/usr/local/lib']
define_macros = [('HAVE_CBLAS', None)]
blis_info:
NOT AVAILABLE
lapack_opt_info:
libraries = ['openblas', 'openblas']
language = c
library_dirs = ['/usr/local/lib']
define_macros = [('HAVE_CBLAS', None)]
blas_opt_info:
libraries = ['openblas', 'openblas']
language = c
library_dirs = ['/usr/local/lib']
define_macros = [('HAVE_CBLAS', None)]
a.flags
:
C_CONTIGUOUS : True
F_CONTIGUOUS : False
OWNDATA : True
WRITEABLE : True
ALIGNED : True
WRITEBACKIFCOPY : False
UPDATEIFCOPY : False
b.flags
:
C_CONTIGUOUS : True
F_CONTIGUOUS : True
OWNDATA : True
WRITEABLE : True
ALIGNED : True
WRITEBACKIFCOPY : False
UPDATEIFCOPY : False
解决方案
我们可以利用multi-core
模块numexpr
来处理大数据并获得内存效率和性能 -
import numexpr as ne
d = {'a0':a[:,0],'b0':b[0],'a1':a[:,1],'b1':b[1],\
'a2':a[:,2],'b2':b[2],'a3':a[:,3],'b3':b[3],\
'a4':a[:,4],'b4':b[4],'a5':a[:,5],'b5':b[5],\
'a6':a[:,6],'b6':b[6],'a7':a[:,7],'b7':b[7]}
eval_str = 'a0*b0 + a1*b1 + a2*b2 + a3*b3 + a4*b4 + a5*b5 + a6*b6 + a7*b7'
out = ne.evaluate(eval_str,d)
计时样本运行 -
In [474]: # Setup with ~10x smaller than posted one, as my system can't handle those
...: np.random.seed(0)
...: a = np.random.randint(0,9,(1000,8,30000)).astype(np.uint8)
...: b = np.random.rand(8).astype(np.float32)
In [478]: %timeit np.einsum('ijk,j->ik', a, b)
1 loop, best of 3: 247 ms per loop
# einsum with optimize flag set as True
In [479]: %timeit np.einsum('ijk,j->ik', a, b, optimize=True)
1 loop, best of 3: 248 ms per loop
In [480]: d = {'a0':a[:,0],'b0':b[0],'a1':a[:,1],'b1':b[1],\
...: 'a2':a[:,2],'b2':b[2],'a3':a[:,3],'b3':b[3],\
...: 'a4':a[:,4],'b4':b[4],'a5':a[:,5],'b5':b[5],\
...: 'a6':a[:,6],'b6':b[6],'a7':a[:,7],'b7':b[7]}
In [481]: eval_str = 'a0*b0 + a1*b1 + a2*b2 + a3*b3 + a4*b4 + a5*b5 + a6*b6 + a7*b7'
In [482]: %timeit ne.evaluate(eval_str,d)
10 loops, best of 3: 94.3 ms per loop
~2.6x
那里的改进。
创建评估部分的更好(不易出错)和通用的方法可能是这样 -
d = {'a'+str(i):a[:,i] for i in range(8)}
d.update({'b'+str(i):b[i] for i in range(8)})
eval_str = ' + '.join(['a'+str(i)+'*'+'b'+str(i) for i in range(8)])
推荐阅读
- swift - 如何修复 iOS 应用程序中的内存泄漏?
- swiftui - Swiftui 隐藏披露箭头
- mongodb - mongoDB updateMany 基于嵌套数组条件
- python - 错误:设置迭代失败:未识别的错误,请检查门户/计算中的日志
- .net - 将文件上传到与 IIS 托管的 ASP.NET 应用程序不同的计算机中的共享文件夹会出现错误,因为对路径“...”的访问被拒绝
- vuejs2 - Vuetify / vuejs 链接在项目列表中不起作用
- scala - 为 Scala 类创建一个空构造函数
- datatables - 如何添加 X 以清除输入字段
- java - 实现远程oauth服务器并与客户端交换令牌而不在服务器上保存令牌
- python - 调用生成器对象返回错误 'TypeError: 'dict' object is not callable'