如何用批次迭代地將ndarray寫入.npy文件

Question

我正在為機器學習應用程序生成大型數據集，它是一個形狀為(N,X,Y)的 numpy 數組。 這里N是樣本數， X是樣本的輸入， Y是樣本的目標。 我想以.npy格式保存這個數組。 我有很多樣本（ N非常大），因此最終數據集大約為 10+ GB。 這意味着我無法創建整個數據集然后保存它，因為它會淹沒我的 memory。

是否可以改為將n樣本的批次迭代寫入該文件？ 因此，我想一次將 append 例如批量 256 個樣本寫入文件（ (256,X,Y) ）。

Answer 1

我發現可以使用np.tofile和np.fromfile 。 請注意，下面的代碼仍然假設您在 memory 中擁有整個數組，但您當然可以更改要動態生成的批次。

import numpy as np

N = 1000;
X = 10;
Y = 1;
my_data = np.random.random((N, X, Y));
print(my_data[700,:,:])

batch_size = 10;

with open('test.dat',mode='wb+') as f:
    i = 0;
    while i < N:
        batch = my_data[i:i+batch_size,:,:]
        batch.tofile(f)

        i += batch_size;

x = np.fromfile('test.dat',dtype=my_data.dtype)

x = np.reshape(x, (N,X,Y))
print(x[700,:,:])

正如@hpaulj 提到的，這個文件不能用np.load加載。

Answer 2

這是一個基於 numpy 的save實現來編寫包含形狀和類型信息的標准npy文件的解決方案：

import numpy as np
import numpy.lib as npl

a = np.random.random((30, 3, 2))
a1 = a[:10]
a2 = a[10:]

filename = 'out.npy'
with open(filename, 'wb+') as f:
    header = npl.format.header_data_from_array_1_0(a1)
    npl.format.write_array_header_1_0(f, header)
    a1.tofile(f)
    a2.tofile(f)
    f.seek(0)
    header['shape'] = (len(a1) + len(a2), *header['shape'][1:])
    npl.format.write_array_header_1_0(f, header)

assert (np.load(filename) == a).all()

這適用於沒有 Python 對象的C_CONTIGUOUS arrays。