Why is my numpy-based custom data loader extremely slow and unstable when iterating over large tick data

I'm currently working on a model similar to DeepLOB, using high-frequency tick-level financial data. Due to the massive volume and the need to structure the data into time series format, it's impossible to load the entire dataset into memory at once.

To handle this, I implemented a custom data loader that reads from .npy files in batches. Here's my Dataloader.py:

import numpy as np
import os
import gc

class Np_DataLoader_Cache:
    
    def __init__(self, data_dir, file_list, batch_size, train_ratio=0.8):
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.file_list = file_list
        self.train_ratio = train_ratio
        
   
        self.current_file_id = -1
        self.current_data = None

       
        self.shapes = []
        for file in file_list:
            with open(os.path.join(data_dir, file), 'rb') as f:
                version = np.lib.format.read_magic(f)
                shape, fortran_order, dtype = np.lib.format.read_array_header_1_0(f)
                self.shapes.append(shape[0])

        self.prefix_sums = [0]
        for s in self.shapes:
            self.prefix_sums.append(self.prefix_sums[-1] + s)

        self.total_samples = self.prefix_sums[-1]
        self.train_end_idx = int(self.total_samples * self.train_ratio)
        self.valid_start_idx = self.train_end_idx

    def _load_file_if_needed(self, file_id):
        
        if self.current_file_id != file_id:
            if self.current_data is not None:
                del self.current_data
                gc.collect()
            filename = self.file_list[file_id]
            self.current_data = np.load(os.path.join(self.data_dir, filename), mmap_mode='r')
            self.current_file_id = file_id

    def _get_batch(self, start_row_global, end_row_global):
        
        files_needed = []
        current_row = start_row_global
        file_id = 0

        while current_row < end_row_global:
            
            while file_id < len(self.prefix_sums) - 1 and self.prefix_sums[file_id + 1] <= current_row:
                file_id += 1

            self._load_file_if_needed(file_id)

            file_start_in_global = self.prefix_sums[file_id]
            file_start_in_file = current_row - file_start_in_global
            file_end_in_global = min(end_row_global, self.prefix_sums[file_id + 1])
            file_end_in_file = file_end_in_global - file_start_in_global

            slice_data = self.current_data[file_start_in_file:file_end_in_file]
            files_needed.append(slice_data)
            current_row += slice_data.shape[0]

        return np.concatenate(files_needed, axis=0)

    def get_train_batch(self, batch_index):
        start_row = batch_index * self.batch_size
        end_row = min((batch_index + 1) * self.batch_size, self.train_end_idx)
        if start_row >= end_row:
            return None
        return self._get_batch(start_row, end_row)

    def get_valid_batch(self, batch_index):
        start_row = self.valid_start_idx + batch_index * self.batch_size
        end_row = min(self.valid_start_idx + (batch_index + 1) * self.batch_size, self.total_samples)
        if start_row >= end_row:
            return None
        return self._get_batch(start_row, end_row)

    def close(self):
        if self.current_data is not None:
            del self.current_data
            self.current_data = None
            self.current_file_id = -1
            gc.collect()

However, during training, I found that the data loading process is extremely slow — not the model computation itself, but specifically the time it takes to iterate over the dataset. To investigate, I ran some experiments to measure the batch loading speed:

import os
data_dir = './LOB_OFI_sortcode_NoResample/'
filelist = sorted([f for f in os.listdir(data_dir) if 'npy' in f])[:20]
from DataLoader_NP_Cache import Np_DataLoader_Cache
demo2 = Np_DataLoader_Cache(data_dir, filelist, batch_size = 4096, train_ratio = 0.8)
batch_nums = demo2.train_end_idx // demo2.batch_size + 1

import time
begin_t = time.time()
very_first = time.time()
for batch_index in range(batch_nums):

    mini_batch = demo2.get_train_batch(batch_index)
    
    if batch_index % 5000 == 0:

        end_t = time.time()
        print( f"Batch of {batch_index} Done. Process Ratio is {batch_index / batch_nums}" )
        elapsed_time = end_t - begin_t
        print(f"This 5000Batchs using time: {elapsed_time:.2f} s")
        begin_t = time.time()

end_t = time.time()
elapsed_time = end_t - very_first
print(f"Total: {elapsed_time:.2f} s")

I noticed that:

Some batches load at a reasonable speed (~15 seconds per 5000 batches),

But occasionally, a batch would suddenly take 100+ seconds, and the speed fluctuates heavily.

These slowdowns appear random — sometimes it recovers, sometimes not. And Sometimes it appears quite early, and sometimes it appears later, just like the picture I posted. But in fact, it often appears for the third or fourth 5000 batchs.

I'm wondering if there's something wrong with how I've implemented the data loader or how numpy is loading data (I'm using np.load(..., mmap_mode='r') internally).

Any suggestions or advice would be greatly appreciated. Thank you!