Computing and loading full dataset into memory (very slow)

[karnawhat]

When preprocessing a dataset that has many structures (>100), the computed molecular features quickly take up lot of available memory in the system (as it has to store 4 different np.ndarray for each example). The system has to constantly free up available memory, which significantly slows down the whole process, resulting in a lot of time to cache the dataset into its respective format.

To combat this, we should not store each of the ndarray's in memory and concatenate them into one big ndarray. We can either:

1. Serialize and save each example as we compute them during the preprocessing step. This might make the code more complex however.
2. Chunk the computed dataset into manageable batches and save each batch.

Similarly, when loading the dataset back into memory, such as when using a InMemorySerializer (see #33), chunk the dataset in batches to reduce memory pressure. Essentially, it should yield an iterator that can properly iterate though each batch. Maybe have it so that we see how much available memory is being used, and if it cannot be fully loaded in memory, then return the iterator instead (with msg saying that we are returning the iterator).

Do a performance test to determine how much chucking the dataset saves both in terms of memory usage and time to process the full dataset.

[karnawhat]

Saving dataset onto disk

For the saving the dataset into disk portion, we will have to:

1. Determine how much free memory is available in the system at the current time to determine the ideal batch_size. Save under pp_utils.py (aka parallel-processing utils)
2. Allow the data to be streamed in batches when preprocessing
   • Specify the batch_size parameter in the parser.parse() function. If a number is provided, it offloads the indices of that size from the dataset to be preprocessed by a thread. If no size is provided, then the ideal size will be determined by the amount of free memory available in the system (see 1. above).
3. Modify the serializers to be able to append information (rather than rewrite) even after the file has been created.
   • Specified via the parameter append=True in serializer.save() method.
   • Allows the info to be saved in parallel right after it is computed.
   • How do we ensure that the keys are stored sequentially rather than by whichever process/thread has its batch completed first? Does the order in which we store each record even matter?
     • Matters: We can directly link the raw protein sequence representation to its preprocessed representation (sequence or 3D mol) through indices.
     • Doesn't: Both 3D structural info (mol) or str (sequence) can be converted back to its original representation through rdkit's MolToSequence() or using AminoAcidTokenizer(), respectively. Additionally, we will likely be shuffling the dataset when training.

This way, once a batch is computed by a thread during the preprocessing step, it can immediately be saved, while the new batch is being computed in a different thread. Obviously, the more processes/threads in the system, the faster it will be.

Loading information from disk

When loading the data from file, we should do the following:

1. Determine the ideal batch_size (see 1. from saving dataset disk)
2. Allow the data to be streamed in batches when loading
   • Only if the whole dataset cannot be loaded into memory. Add warning if this is done (see Warning for loading large datasets #40)
   • Specify the batch_size parameter in serializer.load() function. If a number is provided, it loads the dataset in those batch sizes (as an iterator). If no size is provided, then the ideal size will be determined by the amount of free memory available in the system (see 1. above).