Speeding up Model Training with Multithreading | by Edward Elson Kosasih | Towards Data Science

Waiting for model training to finish could sometimes feel frustrating. Speeding up this process is one of the topmost priority in probably every data scientist’s mind. There are a few approaches that one could try, just to name a few: hardware upgrade (faster CPU/GPU) and model-specific tweaks (e.g. for backpropagation, one can try different optimizers to enable faster convergence). But what if your budget application for hardware upgrade has been pending forever, or none of the existing optimizer bring significant improvement?
Not all hope is lost. We can still make our program more effective. Today, I’m going to talk about multithreading. Specifically, how to concurrently run (1) data loading from disk to RAM (which is CPU-intensive) and (2) model training (which usually is GPU-intensive).
First, let’s load the library dependencies.
import threadingimport timefrom tqdm import tqdm
Assume we have 2 functions: Data loading and Model training. For the sake of this tutorial, I will replace the real function with a time delay method. In this case, Data loading takes 2 seconds while Model training 3 seconds.
# replace with train model - suppose it takes 3 seconds to traindef _trainModel():  time.sleep(3)  return# replace with load data function - suppose it takes 2 seconds to loaddef _loadData():  time.sleep(2)  return
Suppose our model training is in form of backpropagation, which requires many epochs to run. We have chosen an arbitrary number of epochs to loop for: 10 times
# assume we loop 10 timesepochs = 10
If we run the loop as per normal, which is to alternate between data loading and model training, we can see that it takes 50 seconds in total to finish the 10 epochs. From the plot below (which is generated by tqdm), we see that each epoch takes about 5 seconds (2 seconds to load data and 3 seconds to train the model).
original time = data_loading_time + training_model_time
# without multithreading# use with for tqdm to properly shut tqdm down if exception appearswith tqdm(range(epochs)) as epochLoop:  for _ in epochLoop:    # loadData    _loadData()        # trainModel    _trainModel()
Next, we’ll try to run data loading and model training in parallel with python’s threading module. First we have to create a class for each function, extending the threading.Thread parent class.
class TrainModel (threading.Thread):    def __init__(self, data):        threading.Thread.__init__(self)        self.data = data    def run(self):        # return model loss        self._return = _trainModel()        def join(self):        threading.Thread.join(self)        return self._returnclass LoadData (threading.Thread):    def __init__(self, filenames):        threading.Thread.__init__(self)        self.filenames = filenames    def run(self):                # return data        self._return = _loadData()    def join(self):        threading.Thread.join(self)        return self._return
We can then run the functions again by calling the generated thread classes. Here, we can observe that the time taken for 10 epochs has dropped to 30 seconds. That’s a 40% improvement (compared to 50 seconds before)! What happens here is that we’re running both data loading and model training in parallel. The time taken per epoch is no longer a sum of the two functions. It is now the maximum of the two.
multithreaded time = max(data_loading_time, training_model_time)
# with multithreading# use with for tqdm to properly shut tqdm down if exception appearswith tqdm(range(epochs)) as epochLoop:  for _ in epochLoop:    # loadData    loadThread = LoadData(None)    loadThread.start()        # trainModel    trainThread = TrainModel(None)    trainThread.start()        # only continue if both threads are done    modelLoss = trainThread.join()    data = loadThread.join()
We have just seen that multithreading could help to reduce the time taken per epoch in model training. Note that this method will not help much if:
there’s no more CPU resources left for another thread to be spawnedthe original time taken is dominated by a process. If model training takes 10 seconds and data loading 0.1 second, then what normally consumes 101 seconds in total (for 10 epochs) will become 100 seconds. That’s equivalent to just about 1% improvement. In other words, the max of all processes’ time should ideally be way smaller than the their sum in order for multithreading to be effective.
there’s no more CPU resources left for another thread to be spawned
the original time taken is dominated by a process. If model training takes 10 seconds and data loading 0.1 second, then what normally consumes 101 seconds in total (for 10 epochs) will become 100 seconds. That’s equivalent to just about 1% improvement. In other words, the max of all processes’ time should ideally be way smaller than the their sum in order for multithreading to be effective.