Deep Learning with PyTorch
上QQ阅读APP看书,第一时间看更新

Loading PyTorch tensors as batches

It is a common practice in deep learning or machine learning to batch samples of images, as modern graphics processing units (GPUs) and CPUs are optimized to run operations faster on a batch of images. The batch size generally varies depending on the kind of GPU we use. Each GPU has its own memory, which can vary from 2 GB to 12 GB, and sometimes more for commercial GPUs. PyTorch provides the DataLoader class, which takes in a dataset and returns us a batch of images. It abstracts a lot of complexities in batching, such as the usage of multi-workers for applying transformation. The following code converts the previous train and valid datasets into data loaders:

train_data_gen =  
torch.utils.data.DataLoader(train,batch_size=64,num_workers=3)
valid_data_gen =
torch.utils.data.DataLoader(valid,batch_size=64,num_workers=3)

The DataLoader class provides us with a lot of options and some of the most commonly used ones are as follows:

  • shuffle: When true, this shuffles the images every time the data loader is called.
  • num_workers: This is responsible for parallelization. It is common practice to use a number of workers fewer than the number of cores available in your machine.