我想使用dask.delayed
的并行计算来加速函数my_func()
的执行
在三维循环中,my_func()
从iris.cube.Cube
中提取一个值(本质上是从循环外的文件加载的dask.array
),并根据该值,使用networkx
创建一个随机网络,并找到从节点0到节点16的最短路径。每个阵列点的计算是独立的
dask
或multiprocessing
或其他方式来加速它李>以下是一个可复制的示例:
import random
import dask
import iris
import networkx as nx
from dask import delayed
from dask.distributed import Client
from networkx.generators.random_graphs import gnp_random_graph
# Input
fname = iris.sample_data_path("uk_hires.pp") # https://anaconda.org/conda-forge/iris-sample-data
temp_ptntl = iris.load_cube(fname, "air_potential_temperature")[-1, ...] # last time step only
# Dimensions
zs = temp_ptntl.coord("model_level_number").points
lats = temp_ptntl.coord("grid_latitude").points
lons = temp_ptntl.coord("grid_longitude").points
def my_func(iz, iy, ix):
constraint = iris.Constraint(model_level_number=iz, grid_latitude=iy, grid_longitude=ix)
temp_virt = temp_ptntl.extract(constraint) * (1 + 0.61 * 0.04)
if float(temp_virt.data) > 295:
G = nx.gnp_random_graph(30, 0.2, seed=random.randint(1, 10), directed=True)
distance, path = nx.single_source_dijkstra(G, source=0, target=16)
else:
pass
return temp_virt, distance, path
序列号:
%%time
results_serial = [] # serial code
for iz in zs:
for iy in lats[0:5]:
for ix in lons[0:2]:
results_serial.append(my_func(iz, iy, ix))
>>> CPU times: user 2.94 s, sys: 44 ms, total: 2.99 s
>>> Wall time: 2.94 s
使用dask
:
client = Client(processes=True, n_workers=4, threads_per_worker=36)
results_parallel = [] # parallel code
for iz in zs:
for iy in lats[0:5]:
for ix in lons[0:2]:
results_parallel.append(delayed(my_func)(iz, iy, ix))
%%time
computed = dask.compute(results_parallel)
>>> CPU times: user 3.56 s, sys: 344 ms, total: 3.91 s
>>> Wall time: 5.43 s
# client.close()
dask
将有一些开销,因此在小样本上,它表现不佳并不罕见。当我尝试通过更改为for iy in lats[0:15]:
来增加计算量时,我发现串行计算需要10秒,而dask只需要4秒就完成了(还有函数的序列化,这可能需要一些时间,但它仅适用于函数第一次发送到工作程序时)
相关问题 更多 >
编程相关推荐