memory - xarray.Dataset.to_array() 是否将数组加载到内存中以及如何有效地从 xarray 中采样小批量？

Question

我目前正在尝试将一个大的多维数组 (>5 GB) 加载到 python 脚本中。由于我将数组用作机器学习模型的训练数据，因此以小批量高效加载数据非常重要，但要避免将整个数据集加载到内存中一次。
我的想法是使用 xarray图书馆。我使用 X=xarray.open_dataset("Test_file.nc") 加载数据集。据我所知，此命令不会将数据集加载到内存中 - 到目前为止一切正常。但是，我想使用命令 X=X.to_array() 将 X 转换为数组。

我的第一个问题是: X=X.to_array() 是否将它加载到内存中？

如果这样做了，我想知道如何最好地在内存中加载小批量。数组的形状是(variable,datetime,x1_position,x2_position)。我想按日期时间加载小批量，这将导致:

ind=np.random.randint(low=0,high=n_times,size=(BATCH_SIZE))
mini_batch=X[:,ind]

另一种方法是在使用 X.transpose("datetime","variable","x1_position","x2_position") 之前转置数组，然后通过以下方式采样:

ind=np.random.randint(low=0,high=n_times,size=(BATCH_SIZE))
mini_batch=X[ind,:]

我的第二个问题是:转置 xarray 会影响索引效率吗？更具体地说，X[ind,:] 是否与 X[:,ind] 一样长？

Answer 1

My first question is: Does X=X.to_array() load it into memory or not?

xarray 利用 dask 将部分数据分块(加载)到内存中。你可以比较X通过

X = xarray.open_dataset("Test_file.nc")
# or
X = xarray.open_dataset("Test_file.nc",
         chunks={'datetime':1, 'x1_position':x1_count, 'x2_position':x2_count})

并查看(print(X))加载的数据集之间的差异，或相应地指定 block 。

后一种方式意味着仅将一个datetime 切片数据分块(加载)到内存中。我认为您不需要 X=X.to_array()，但您也可以比较 to_array() 之后的结果。我的经验是 to_array() 不会改变实际的分块(加载)，而只会改变数据的 View 。

My second question is: Does transposing an xarray affect the efficiency of indexing? More specifically, does X[ind,:] take as long as X[:,ind]?

我认为xarray的一个目标是让用户忘记底层实现的细节(基于numpy)。转置可能只会修改 View 而不是数据的底层结构。两种索引方式之间肯定存在一些效率差异，具体取决于哪种方式沿着连续内存访问数据。但这种差异不会是开销。随意使用两者。