Load data into GPU directly using PyTorch

user_13606134 0
Viewed 16893

In training loop, I load a batch of data into CPU and then transfer it to GPU:

import torch.utils as utils

train_loader = utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4, pin_memory=True)

for inputs, labels in train_loader:
    inputs, labels = inputs.to(device), labels.to(device)

This way of loading data is very time-consuming. Any way to directly load data into GPU without transfer step ?

dataloader pytorch deep-learning gpu python-3.x
2 Answers

@PeterJulian first of all thanks for the reply. As far as I know there is no single line command for loading a whole dataset to GPU. Actually in my reply I meant to use .to(device) in the __init__ of the data loader. There are some examples in the link that I had shared previously. Also, I left an example data loader code below. Hope both the examples in the link and the code below helps.

class SampleDataset(Dataset):
    def __init__(self, device='cuda'):
        super(SampleDataset, self).__init__()
        self.data = torch.ones(1000)
        self.data = self.data.to(device)
    
    def __len__(self):
        return len(self.data)

    def __getitem__(self, i):
        element = self.data[i]
        return element
user_128709640

You can load all the data to in tensor than move it yo GPU memory.(assuming that you have enough memory) When you need it use the one inside the tensor which is already at GPU memory. Hope it helps.

user_128709640
Related