I created a training loop with pytorch's TensorDataset and DataLoader classes, but encounter an interesting behavior. The progress intermittently halts every 10-15 batches with seemingly no reason. I tried adding shared memory, increasing the num_workers parameter and modified some other parameters of the dataloader, but could not resolve the issue. Below you can find a reproducible example: X_train = torch.from_numpy(X_train).share_memory_() # float32 dtype, shape=(10_000_000, 400) y_train = torch.from_numpy(y_train).share_memory_() # float32 dtype, (10_000_000, ) train_dataset = torch.utils.data.TensorDataset(X_train, y_train) train_dataloader = torch.utils.data.DataLoader( train_dataset, batch_size = 8192, shuffle = True, num_workers = 4, pin_memory = True, prefetch_factor = 2, persistent_workers = True, ) ... # training loop epoch_st_time = time.time() for i, (xt, yt) in enumerate(tqdm(train_dataloader, desc="Training...")): loop_st_time = time.time() xt = xt.to(device) yt = yt.to(device) # other training code, issue persists even after removal of them ... loop_ed_time = time.time() epoch_ed_time = time.time() I also measured the time it takes for code to execute inside the loop (loop_ed_time-loop_st_time) where it takes at most 0.009 and on average 0.0001 seconds to complete one loop. When I take the sum the times it took inside the loops it measures up to 1.5 seconds whilst the time it takes to complete the epoch (epoch_ed_time-epoch_st_time) is 72 seconds. Part of this is the time it takes to start the loop where I guess pytorch creates parallel processes and memory allocations for multiple workers (loop_st_time-epoch_st_time at i=0), but this is 8 seconds. I think the cause for this long execution time and intermittent halts is something related with the data collation process of the dataloader, but I could not resolve the issue. For comparison, using a keras model and calling a fit function takes around 10 seconds for one epoch to complete. Continue reading...
