Loading from a huge HDF5 dataset with jax.random.choice

[nalzok]

I would like to sample from the 3d-shapes dataset with some weight on each individual specimen. Here is the full source code, and you need to download the 255.2 MB dataset from Google Cloud Storage and place the file in data/3dshapes.h5.

In the dataset, data['images'] and data['labels'] have size (480000, 64, 64, 3) and (480000, 6) respectively, but I am only interested in one of the labels, hence Y = jnp.array(data['labels'][:, label_col]).

# dataloader.py

import h5py

import jax
import jax.numpy as jnp

DATA_PATH = 'data/3dshapes.h5'

# floor hue:    10 values linearly spaced in [0, 1]
# wall hue:     10 values linearly spaced in [0, 1]
# object hue:   10 values linearly spaced in [0, 1]
# scale:        8 values linearly spaced in [0, 1]
# shape:        4 values in [0, 1, 2, 3]
# orientation:  15 values linearly spaced in [-30, 30]
num_classes = [10, 10, 10, 8, 4, 15]


def dataloader(weight, label_col, batch_size, *, key):
    with h5py.File(DATA_PATH) as data:
        # Treat these as the population, not a sample
        X = jnp.array(data['images'], dtype=jnp.int8)   # NHWC
        Y = jnp.array(data['labels'][:, label_col])

    @jax.jit
    def next_batch(key):
        key, data_key = jax.random.split(key)
        idx = jax.random.choice(data_key, X.shape[0], (batch_size,), p=weight)
        return key, X[idx]/255., Y[idx]

    while True:
        print('getting next batch...')
        key, X_batch, Y_batch = next_batch(key)
        print('got next batch')
        yield X_batch, Y_batch

I tried to load one batch to see if it's working

if __name__ == '__main__':
    weight = jnp.ones(480000) / 480000
    label_col = 4
    batch_size = 8

    loader_rng = jax.random.PRNGKey(42)
    loader = dataloader(weight, label_col, batch_size, key=loader_rng)
    X, y = next(loader)

This turns out to be extremely slow, and I got the following error after several minutes. Note that each batch merely has size (8, 64, 64, 3), and I am on a v3-8 TPU VM with 335GB RAM, so the available memory should be more than enough.

$ pipenv run python dataloader.py
tcmalloc: large alloc 5898240000 bytes == 0x55a8d17c0000 @  0x7f0dbb515680 0x7f0dbb536824 0x7f0db960afa4 0x7f0db960b6bf 0x7f0db966c356 0x7f0db966e645 0x7f0db9708b5b 0x55a84bdbb66a 0x55a84bbe11b7 0x55a84bcaf198 0x55a84bbf664c 0x7f0d185beb48 0x55a84bbf5c2c 0x55a84bda0b72 0x55a84bbf62f3 0x55a84bbf766d 0x7f0db966773f 0x7f0db966d512 0x7f0db964bedd 0x7f0db964ce91 0x7f0db966d5f9 0x7f0db966dbb7 0x7f0db970af91 0x55a84bdbb66a 0x55a84bbe11b7 0x55a84bcaf198 0x55a84bbe11b7 0x55a84bdae7bd 0x55a84bdf4c77 0x55a84bdbb7dc 0x55a84bbe172e
tcmalloc: large alloc 5898240000 bytes == 0x55aa38c9c000 @  0x7f0dbb515680 0x7f0dbb536824 0x7f0dbb536b8a 0x7f0c71830afc 0x7f0c6f1b4a1d 0x7f0c6d219fec 0x7f0c6d216992 0x7f0c6d21d0fd 0x7f0c6d20d1f1 0x7f0c6cfe8b90 0x7f0c6cfd49e6 0x55a84bdbb033 0x55a84bbf5c2c 0x55a84bda0b52 0x55a84bbe1ece 0x55a84bcaf198 0x55a84bbe172e 0x55a84bcaf198 0x55a84bbe172e 0x55a84bcaf198 0x55a84bbe172e 0x55a84bcaf198 0x55a84bd4c24d 0x55a84bbf664c 0x55a84bbdaaa2 0x55a84bcaf198 0x55a84bbf664c 0x55a84bbdaaa2 0x55a84bcaf198 0x55a84bd4c313 0x55a84bbf66a6
getting next batch...
tcmalloc: large alloc 5898248192 bytes == 0x55ab98fb2000 @  0x7f0dbb515680 0x7f0dbb536824 0x7f0dbb536b8a 0x7f0c71830afc 0x7f0c71613b98 0x7f0c71613e49 0x7f0c716141df 0x7f0c6d1ff7d3 0x7f0c6d2039a3 0x7f0c6cfd49e6 0x55a84bdbb033 0x55a84bbf5c2c 0x55a84bda0b52 0x55a84bbe1ece 0x55a84bcaf198 0x55a84bdf7094 0x55a84bc0aa0d 0x55a84bc0ad6f 0x55a84bc0af00 0x55a84bbe172e 0x55a84bcaf198 0x55a84bbe172e 0x55a84bcaf198 0x55a84bbf664c 0x55a84bbdaaa2 0x55a84bcaf198 0x55a84bbf664c 0x55a84bbdaaa2 0x55a84bcaf198 0x55a84bbf664c 0x55a84bbdaaa2
tcmalloc: large alloc 5898240000 bytes == 0x55acf90b4000 @  0x7f0dbb515680 0x7f0dbb536824 0x7f0c64e0bcd1 0x7f0c64e0c3d1 0x7f0c67fdc5d3 0x7f0c67fdd09a 0x7f0c64e210c0 0x7f0c6ae74195 0x7f0c6ae6fec8 0x7f0c6ae69c85 0x7f0c64e1d430 0x7f0c64e0f035 0x7f0c70dad9b0 0x7f0c6f1b0cd1 0x7f0c7182dd95 0x7f0c7182be93 0x7f0c7181373b 0x7f0dbb4d9609 0x7f0dbb2a4133
tcmalloc: large alloc 5898240000 bytes == 0x55ae589b4000 @  0x7f0dbb515680 0x7f0dbb536824 0x7f0dbb536b8a 0x7f0c64e22f09 0x7f0c64e22098 0x7f0c6ae74195 0x7f0c6ae6fec8 0x7f0c6ae69c85 0x7f0c64e1dbb4 0x7f0c64e0f1a7 0x7f0c70dadd7d 0x7f0c6f1bdbb5 0x7f0c6d1ff86f 0x7f0c6d2039a3 0x7f0c6cfd49e6 0x55a84bdbb033 0x55a84bbf5c2c 0x55a84bda0b52 0x55a84bbe1ece 0x55a84bcaf198 0x55a84bdf7094 0x55a84bc0aa0d 0x55a84bc0ad6f 0x55a84bc0af00 0x55a84bbe172e 0x55a84bcaf198 0x55a84bbe172e 0x55a84bcaf198 0x55a84bbf664c 0x55a84bbdaaa2 0x55a84bcaf198
tcmalloc: large alloc 5898240000 bytes == 0x55acf90b4000 @  0x7f0dbb515680 0x7f0dbb536824 0x7f0c67ff8a3b 0x7f0c67ff602b 0x7f0c64e22fe5 0x7f0c6b639aaf 0x7f0c681b19fb 0x7f0c681bbf01 0x7f0c67a6268f 0x7f0c681c2ecb 0x7f0c6b8bbbaa 0x7f0c6b8c174b 0x7f0c6b8be265 0x7f0c6bb2341e 0x7f0dbb4d9609 0x7f0dbb2a4133
tcmalloc: large alloc 5898240000 bytes == 0x55aa38bb2000 @  0x7f0dbb515680 0x7f0dbb536824 0x7f0db960afa4 0x7f0db960b6bf 0x7f0db966c356 0x7f0db966c447 0x7f0db965e0e6 0x7f0db970d2a0 0x55a84bdbb66a 0x55a84bbe1ece 0x55a84bcaf198 0x55a84bbe172e 0x55a84bcaf198 0x55a84bbe172e 0x55a84bcaf198 0x55a84bbe172e 0x55a84bcaf198 0x55a84bdf7094 0x55a84bc0aa0d 0x55a84bc0ad6f 0x55a84bc0af00 0x55a84bbe172e 0x55a84bcaf198 0x55a84bbe172e 0x55a84bcaf198 0x55a84bbe11b7 0x55a84bcaf198 0x55a84bbe1ece 0x55a84bcaf198 0x55a84bbf664c 0x55a84bbdaaa2
tcmalloc: large alloc 5898248192 bytes == 0x55acf90b4000 @  0x7f0dbb515680 0x7f0dbb535ff4 0x7f0c54773dfd 0x7f0c5541f89c 0x7f0c554a95a7 0x7f0c55422462 0x7f0c5530cb7c 0x7f0c55a81d62 0x7f0c55a98a39 0x7f0c55a53a94 0x55a84bdbb033 0x55a84bbf5c2c 0x55a84bbe17e6 0x55a84bcaf198 0x55a84bbe172e 0x55a84bcaf198 0x55a84bbe172e 0x55a84bcaf198 0x55a84bbe172e 0x55a84bcaf198 0x55a84bdf7094 0x55a84bc0aa0d 0x55a84bc0ad6f 0x55a84bc0af00 0x55a84bbe172e 0x55a84bcaf198 0x55a84bbe172e 0x55a84bcaf198 0x55a84bbe11b7 0x55a84bcaf198 0x55a84bbe1ece
tcmalloc: large alloc 11796480000 bytes == 0x55afb8ab4000 @  0x7f0dbb515680 0x7f0dbb536824 0x7f0c554d5a50 0x7f0c553fa837 0x7f0c553fba69 0x7f0c55401c5a 0x7f0c5502203a 0x7f0c55022086 0x7f0c55409197 0x7f0c55409a5b 0x7f0c5540a21b 0x7f0c5546b5e6 0x7f0c5531bfab 0x7f0c5531c2a6 0x7f0c55409197 0x7f0c55409a5b 0x7f0c5540a21b 0x7f0c5544096d 0x7f0c554409c6 0x7f0c55409197 0x7f0c5540a5ba 0x7f0c5540f2cb 0x7f0c55314a78 0x7f0c55aa0ea5 0x7f0c55ad57cc 0x7f0c55a53a94 0x55a84bdbb033 0x55a84bbf5c2c 0x55a84bda0b52 0x55a84bbe11b7 0x55a84bcaf198
tcmalloc: large alloc 11796488192 bytes == 0x55b278cb4000 @  0x7f0dbb515680 0x7f0dbb535ff4 0x7f0c553fa7be 0x7f0c553fba69 0x7f0c55401c5a 0x7f0c5502203a 0x7f0c55022086 0x7f0c55409197 0x7f0c55409a5b 0x7f0c5540a21b 0x7f0c5546b5e6 0x7f0c5531bfab 0x7f0c5531c2a6 0x7f0c55409197 0x7f0c55409a5b 0x7f0c5540a21b 0x7f0c5544096d 0x7f0c554409c6 0x7f0c55409197 0x7f0c5540a5ba 0x7f0c5540f2cb 0x7f0c55314a78 0x7f0c55aa0ea5 0x7f0c55ad57cc 0x7f0c55a53a94 0x55a84bdbb033 0x55a84bbf5c2c 0x55a84bda0b52 0x55a84bbe11b7 0x55a84bcaf198 0x55a84bbe1ece
tcmalloc: large alloc 11796488192 bytes == 0x55b5386b6000 @  0x7f0dbb515680 0x7f0dbb536824 0x55a84bc368b5 0x55a84bc66600 0x55a84bc93048 0x7f0c55aa7c75 0x7f0c5530ab27 0x7f0c554e56a6 0x7f0c553fba79 0x7f0c55401c5a 0x7f0c5502203a 0x7f0c55022086 0x7f0c55409197 0x7f0c55409a5b 0x7f0c5540a21b 0x7f0c5546b5e6 0x7f0c5531bfab 0x7f0c5531c2a6 0x7f0c55409197 0x7f0c55409a5b 0x7f0c5540a21b 0x7f0c5544096d 0x7f0c554409c6 0x7f0c55409197 0x7f0c5540a5ba 0x7f0c5540f2cb 0x7f0c55314a78 0x7f0c55aa0ea5 0x7f0c55ad57cc 0x7f0c55a53a94 0x55a84bdbb033
tcmalloc: large alloc 11800395776 bytes == 0x55b7f80b8000 @  0x7f0dbb515680 0x7f0dbb536824 0x55a84bc368b5 0x55a84bc66600 0x55a84bc8b7bf 0x55a84bc8ba62 0x55a84bd8ed48 0x55a84bd7a53c 0x55a84bdaa058 0x55a84bbe01d6 0x55a84bcaf198 0x55a84bbe1ece 0x55a84bcaf198 0x55a84bbe172e 0x55a84bcaf198 0x55a84bbf66a6 0x55a84bbdaaa2 0x55a84bcaf198 0x55a84bbe01d6 0x55a84bcaf198 0x55a84bbf664c 0x55a84bbdaaa2 0x55a84bcaf198 0x55a84bbf664c 0x55a84bbdaaa2 0x55a84bcaf198 0x55a84bbf66a6 0x55a84bbdaaa2 0x55a84bcaf198 0x55a84bbe01d6 0x55a84bcaf198
tcmalloc: large alloc 11800395776 bytes == 0x55bab8674000 @  0x7f0dbb515680 0x7f0dbb535ff4 0x7f0c6cfbd07e 0x7f0c6cfbf819 0x7f0c6cff5593 0x7f0c6cfd49e6 0x55a84bdbb033 0x55a84bbf5c2c 0x55a84bda0b52 0x55a84bbe11b7 0x55a84bcaf198 0x55a84bbf664c 0x55a84bbdaaa2 0x55a84bcaf198 0x55a84bbe172e 0x55a84bcaf198 0x55a84bbe172e 0x55a84bcaf198 0x55a84bbf66a6 0x55a84bbdaaa2 0x55a84bcaf198 0x55a84bbe01d6 0x55a84bcaf198 0x55a84bbf664c 0x55a84bbdaaa2 0x55a84bcaf198 0x55a84bbf664c 0x55a84bbdaaa2 0x55a84bcaf198 0x55a84bbf66a6 0x55a84bbdaaa2
tcmalloc: large alloc 5898248192 bytes == 0x55a8d173e000 @  0x7f0dbb515680 0x7f0dbb535ff4 0x7f0c71b7e34d 0x7f0c71b7e860 0x7f0c70e1f10e 0x7f0c70e2d2ea 0x7f0c70e31069 0x7f0c70e33011 0x7f0c70e34015 0x7f0c70e0ebf8 0x7f0c6fdb2c61 0x7f0c70e1b6a5 0x7f0c70e1d807 0x7f0c70e1d8f8 0x7f0c70e1e0e7 0x7f0c70e1e79e 0x7f0c70e1e963 0x7f0c7142c889 0x7f0c6ffa344a 0x7f0c70e1b6a5 0x7f0c70e1d807 0x7f0c70e1d8f8 0x7f0c70e1e0e7 0x7f0c70e1e79e 0x7f0c70e1e831 0x7f0c713fff6c 0x7f0c70e1b6a5 0x7f0c70e1cccc 0x7f0c70e1d70b 0x7f0c6f1c0791 0x7f0c6d210223
tcmalloc: large alloc 5898248192 bytes == 0x55aa38bb2000 @  0x7f0dbb515680 0x7f0dbb535ff4 0x7f0c6e09bc0d 0x7f0c713e4b2c 0x7f0c71467cc7 0x7f0c713e7472 0x7f0c70e31fe5 0x7f0c70e33011 0x7f0c70e34015 0x7f0c70e0ebf8 0x7f0c6fdb2c61 0x7f0c70e1b6a5 0x7f0c70e1d807 0x7f0c70e1d8f8 0x7f0c70e1e0e7 0x7f0c70e1e79e 0x7f0c70e1e963 0x7f0c7142c889 0x7f0c6ffa344a 0x7f0c70e1b6a5 0x7f0c70e1d807 0x7f0c70e1d8f8 0x7f0c70e1e0e7 0x7f0c70e1e79e 0x7f0c70e1e831 0x7f0c713fff6c 0x7f0c70e1b6a5 0x7f0c70e1cccc 0x7f0c70e1d70b 0x7f0c6f1c0791 0x7f0c6d210223
[libprotobuf ERROR external/com_google_protobuf/src/google/protobuf/message_lite.cc:484] xla.HloModuleProto exceeded maximum protobuf size of 2GB: 5900227430
2022-05-21 14:25:03.090270: F external/org_tensorflow/tensorflow/stream_executor/tpu/proto_helper.h:44] Check failed: proto.SerializeToArray(bytes, size)
https://symbolize.stripped_domain/r/?trace=7f0dbb1c800b,7f0dbb4e541f,7f0c6ce94132,7f0c6ce94515,7f0c70079426,7f0c7006f586,7f0c700689b8,7f0c6f1afd80,7f0c6f1be23f,7f0c6d2105a9,7f0c6cff4fb5,7f0c6cff55f6,7f0c6cfd49e5,55a84bdbb032,55a84bf52c1f&map=
*** SIGABRT received by PID 1252966 (TID 1252966) on cpu 50 from PID 1252966; stack trace: ***
PC: @     0x7f0dbb1c800b  (unknown)  raise
    @     0x7f0c6bc15c73        992  (unknown)
    @     0x7f0dbb4e5420  (unknown)  (unknown)
    @     0x7f0c6ce94133        432  stream_executor::tpu::SerializeProto<>()
    @     0x7f0c6ce94516       3440  xla::(anonymous namespace)::TpuCompiler::RunHloPasses()
    @     0x7f0c70079427        608  xla::Service::BuildExecutable()
    @     0x7f0c7006f587       1152  xla::LocalService::CompileExecutables()
    @     0x7f0c700689b9       2672  xla::LocalClient::Compile()
    @     0x7f0c6f1afd81        896  xla::PjRtStreamExecutorClient::Compile()
    @     0x7f0c6f1be240       1216  xla::PjRtStreamExecutorClient::Compile()
    @     0x7f0c6d2105aa       1168  xla::PyClient::CompileMlir()
    @     0x7f0c6cff4fb6       1888  pybind11::detail::argument_loader<>::call_impl<>()
    @     0x7f0c6cff55f7        208  pybind11::cpp_function::initialize<>()::{lambda()#3}::_FUN()
    @     0x7f0c6cfd49e6        768  pybind11::cpp_function::dispatcher()
    @     0x55a84bdbb033  (unknown)  cfunction_call
    @     0x55a84bf52c20  (unknown)  (unknown)
https://symbolize.stripped_domain/r/?trace=7f0dbb1c800b,7f0c6bc15c72,7f0dbb4e541f,7f0c6ce94132,7f0c6ce94515,7f0c70079426,7f0c7006f586,7f0c700689b8,7f0c6f1afd80,7f0c6f1be23f,7f0c6d2105a9,7f0c6cff4fb5,7f0c6cff55f6,7f0c6cfd49e5,55a84bdbb032,55a84bf52c1f&map=abc33f1bfca16f4e7d925d4248b4beb3:7f0c574e7000-7f0c6bf96b70
E0521 14:25:03.177025 1252966 coredump_hook.cc:366] RAW: Remote crash data gathering hook invoked.
E0521 14:25:03.177085 1252966 coredump_hook.cc:412] RAW: Skipping coredump since rlimit was 0 at process start.
E0521 14:25:03.177097 1252966 client.cc:234] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
E0521 14:25:03.177107 1252966 coredump_hook.cc:473] RAW: Sending fingerprint to remote end.
E0521 14:25:03.177120 1252966 coredump_socket.cc:118] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket
E0521 14:25:03.177132 1252966 coredump_hook.cc:477] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?
E0521 14:25:03.177141 1252966 coredump_hook.cc:551] RAW: Discarding core.
E0521 14:25:03.447653 1252966 process_state.cc:774] RAW: Raising signal 6 with default behavior
fish: “pipenv run python dataloader.py” terminated by signal SIGABRT (Abort)

[jakevdp]

JAX doesn't have any mechanism to only load parts of an array from disk – the code you wrote will attempt to load the entire dataset, and only then pull out the subset of relevant data. So, for example, when you call

X = jnp.array(data['images'], dtype=jnp.int8)

I believe that JAX will convert data['images'] to a jnp.ndarray, which pulls all the data into memory.

If you want to use h5py to load the data in batches, you'll have to keep data['images'] as an h5py object, and avoid attempting to convert it to a jax array. And since jit only works with JAX objects, you'll want to avoid JIT-compiling your next_batch function. At that point, I think your approach should work correctly.

[nalzok]

Yeah, so I actually don't want to do lazy loading at all. As you can see in the following snippet, I tried to eagerly read the whole dataset into a jax.numpy.ndarray.

with h5py.File(DATA_PATH) as data:
    # Treat these as the population, not a sample
    X = jnp.array(data['images'], dtype=jnp.int8)   # NHWC
    Y = jnp.array(data['labels'][:, label_col])

That way, when taking a slice e.g. X[idx]/255., Y[idx], I can just look at the data pre-loaded into the fast device memory and don't need to access the disk anymore. I suppose that will give me better performance.

[jakevdp]

Understood - I’d assumed that step was leading to the resource exhaustion, because you load a multi-gigabyte dataset entirely into memory. Is that not the case?

[nalzok]

I don't think so. If you look closely at the error message, there is a line getting next batch.... That print statement won't be executed after the array has been fully loaded into the memory, so the error must have occurred somewhere inside next_batch.

[jakevdp]

OK, I see – I was thrown off by your emphasis on HDF5 and batching in the question, but now I see the question doesn't really have anything to do with HDF5, except indicentally in that that's how you created the initial array (had you used jnp.empty((480000, 64, 64, 3), dtype='uint8') I suspect you'd see the same result). Unfortunately, I don't have any good ideas about what might be causing multiple allocations of copies of your input array.

[nalzok]

That’s fine. Thanks for your time and patience anyway!