Home > Software engineering >  When to put pytorch tensor on GPU?
When to put pytorch tensor on GPU?

Time:10-13

I'm experimenting with running neural network on GPU using pytorch, and my data have some unusual shape so I use Dataset and DataLoader to generate data batch. My code runs fine on CPU but I'm a little confused on when is the right timing to put the data on GPU:

  1. My data size is small enough to be put all together on GPU, should I put all data on GPU before fitting, so that all DataLoader and Dataset operations only take place on GPU in order to get optimal execution speed?
  2. Another possibility is to leave all data on CPU which could be useful when the data size become larger. In that case, should I call batch.to("cuda") for each batch generated from DataLoader?
  3. Should I also put the model on GPU first before training? It is a small enough model to be put on GPU.
  4. My raw data are numpy array, hence I have the freedom to write Dataset that returns numpy array in __getitem()___ method, or convert the numpy array to pytorch tensor and write Dataset that returns pytorch tensor. Is one method preferred over the other?

CodePudding user response:

If you are looking to use a GPU device for training a PyTorch model, you should:

  • 1. and 2. Place your model on the GPU, it will stay there for the duration of the training.
  • 3. and 4. Leave both the dataset and data loader processing on the CPU. If time you fetch a batch, your dataloader will request some instances from the dataset and return them. You will have to transfer them to the appropriate device: here the GPU to perform inference. Do note, inputs and the model need to be on the same device. Lastly, the Numpy array to PyTorch tensor will be handled by the data loader so your __getitem__ can return Numpy arrays.

CodePudding user response:

Let me clear up a thing. at the time of passing the data through the model, both your model and the data(that specific batch) have to be on the same device. To automate your code to work on both GPU and non-GPU environments you might use this line.

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

So if You are looking to using GPU for training you must put the model on GPU before training. personally, I prefer loading the model to GPU whenever I am creating the model object. [Answer 3]

model = MyModel().to(device)

Then you also need to put your data in GPU. One obvious option is putting all data at once. But I would suggest against that. Because however small your dataset is, You will always get better performance by putting one batch at a time in GPU, rather than whole data at once. I know you are thinking, there is a time delay for taking each batch from CPU to GPU. You are right!

But putting one batch at a time would allow you a larger batch size, and a larger batch size will always win in terms of performance, compared to smaller batch sizes and loading all data at a time.[Answer 1 & 2]

for x_data, y_data in train_dataloader:
    x_data, y_data = x_data.to(device), y_data.to(device)

Finally, about writing the __getitem__, the Numpy array to PyTorch tensor will be handled by the data loader so your getitem can return Numpy arrays. But I feel good when I see the conversion written explicitly in my code. It gives me the sense of a complete and easy-to-understand pipeline.[Answer 4]

  • Related