In the examples below, a 500x2000x2000 three-dimensional NumPy array named a is initialized with zeros. At each iteration, a random two-dimensional array r is inserted into array a. This approach could represent a larger piece of code where the r array would be created from various calculations during each iteration of the for-loop. Consequently, each slice in the z dimension of the a array is calculated at each iteration.
The examples were run on a MacBook Pro with a 2.6 GHz 6-Core Intel Core i7 CPU and 32 GB of memory (RAM). Memory usage is profiled with the memory-profiler tool using the terminal commands shown below.
# Generate a memory profile for ex1_basic.py
mprof run --output ex1_mprof.dat ex1_basic.py
mprof plot --output ex1_mprof.pdf ex1_mprof.dat
Example 1
Basic example of building a large NumPy array with random values. Notice that np.random.rand() creates an array of np.float64 values.
# ex1_basic.py
import numpy as np
import time
def main():
tic = time.perf_counter()
z = 500 # depth
x = 2000 # rows
y = 2000 # columns
a = np.zeros((z, x, y))
for i in range(z):
r = np.random.rand(x, y)
a[i] = r
toc = time.perf_counter()
print('elapsed time =', round(toc - tic, 2), 'sec')
s = np.float64().nbytes * (z * x * y) / 1e9 # where 1 GB = 1000 MB
print('calculated storage =', s, 'GB')
if __name__ == '__main__':
main()
Using the h5py package, this example creates an hdf5 file that contains a dataset that represents the a array. The dset variable is similar to the a variable in the previous examples. This approach allows the array to reside on disk, not in memory.
# ex3_hdf5.py
import numpy as np
import h5py
import time
def main():
rng = np.random.default_rng()
tic = time.perf_counter()
z = 500 # depth
x = 2000 # rows
y = 2000 # columns
f = h5py.File('file.hdf5', 'w')
dset = f.create_dataset('data', shape=(z, x, y), dtype=np.float32)
for i in range(z):
r = rng.standard_normal((x, y), dtype=np.float32)
dset[i, :, :] = r
toc = time.perf_counter()
print('elapsed time =', round(toc - tic, 2), 'sec')
s = np.float32().nbytes * (z * x * y) / 1e9 # where 1 GB = 1000 MB
print('calculated storage =', s, 'GB')
if __name__ == '__main__':
main()
A comparison of the three examples is shown below in the table. Changing the array data type from float64 to float32 cut the memory usage in half. Writing the array to disk drastically reduced memory use. The generated hdf5 file is about 8 GB on disk which is the size of the array containing float32 values. The elapsed time for the hdf5 approach is similar to the other examples; therefore, writing to the hdf5 file seems to have a negligible performance impact. Based on these examples, writing the array to disk using an hdf5 file allows substantial memory reduction as long as disk space is available for storing the array.
Example
Elapsed time
Calculated storage
Peak memory
File storage
1 basic
16.56 s
16.0 GB
16.05 GB
n/a
2 float32
19.41 s
8.0 GB
8.06 GB
n/a
3 hdf5
21.34 s
8.0 GB
104 MB
8.02 GB
Other approaches for writing the array to disk such as numpy.memmap and zarr were tested too. But the hdf5 approach demonstrated in example 3 gave better performance and memory usage for this type of problem.
A comment about the memory increase over time for examples 1 and 2. When an array is initialized with np.zeros, the memory is lazily allocated as values are added to the array. The total amount of memory allocated for the array is based on the data type of the values stored in the array.
