使用xarray dask apply_ufunc时避免内存溢出

Question

我需要沿着以下形状的xarray-dask数组的时间维度应用一个函数：

<xarray.DataArray 'tasmax' (time: 14395, lat: 1801, lon: 3600)>
dask.array<rechunk-merge, shape=(14395, 1801, 3600), dtype=float32, chunksize=(14395, 1801, 600), chunktype=numpy.ndarray>
Coordinates:
* lat      (lat) float64 90.0 89.9 89.8 89.7 89.6 ... -89.7 -89.8 -89.9 -90.0
* time     (time) datetime64[ns] 1981-01-01 1981-01-02 ... 2020-05-31
* lon      (lon) float64 -180.0 -179.9 -179.8 -179.7 ... 179.7 179.8 179.9

该过程的输出将是一个小得多的数组，大小为(time=365, lat=1801, lon=3600)，但如上所示，输入数组内存大小约为360GB。我有一台有16个CPU核心和126 GB RAM的机器。我试图通过使用带有dask='parallelized'参数的apply_ufunc来优化进程，但由于使用了所有126GB RAM，这会导致内存错误。我可以避免并行化，但完成这个过程需要很长时间。有没有什么方法可以控制apply_ufunc的内存使用，将进程包含在我可用的RAM中？

Answer 1

您使用的chunksize太大。正如Val在评论中提到的，建议每个chunck的大小约为100mb。

您可以将1801更改为1-4之间的值以获得更好的性能：

chunksize=(14395, 4, 600)