Optimizing PyTorch training by wrapping torch.utils.data.Dataset with tensordict.TensorDict.MemoryMappedTensor mapped, pinned,
and loaded onto an Nvidia GPU and inputting TensorDict(Dataset) into torch.utils.data.DataLoader--to boost model training speed.
To run the demo:
git clone https://github.com/OriYarden/pytorch_training_optimization_using_tensordict_memory_mapping
cd pytorch_training_optimization_using_tensordict_memory_mapping
python run_demo.py
Demo (as gifs):
torch.utils.data.Dataset # Training 1 Epoch:
tensordict.TensorDict.MemoryMappedTensor(torch.utils.data.Dataset) # Training 1 Epoch:
TensorDict Memory Mapping boosts training speed.
The initial wrapping runtime is approximately equal to 1 epoch of torch.utils.data.Dataset:
