The Most Complete Guide to PyTorch for Data Scientists
By: Rahul Agarwal
Tags: Programming, Artificial Intelligence, Machine Learning, Data Science, Deep Learning

Understanding Custom DataLoaders

This particular section is a little advanced and can be skipped going through this post as it will not be needed in a lot of situations. But I am adding it for completeness here.

So let’s say you are looking to provide batches to a network that processes text input, and the network could take sequences with any sequence size as long as the size remains constant in the batch. For example, we can have a BiLSTM network that can process sequences of any length. It’s alright if you don’t understand the layers used in it right now; just know that it can process sequences with variable sizes.

This network expects its input to be of shape ( batch_size , seq_length ) and works with any seq_length . We can check this by passing our model two random batches with different sequence lengths(10 and 25).

model = BiLSTM()

input_batch_1 = torch.randint(low = 0,high = 10000, size = (100,10))

input_batch_2 = torch.randint(low = 0,high = 10000, size = (100,25))

print(model(input_batch_1).size())

print(model(input_batch_2).size())

------------------------------------------------------------------

torch.Size([100, 1])

torch.Size([100, 1])

Now, we want to provide tight batches to this model, such that each batch has the same sequence length based on the max sequence length in the batch to minimize padding. This has an added benefit of making the neural net run faster. It was, in fact, one of the methods used in the winning submission of the Quora Insincere challenge in Kaggle, where running time was of utmost importance.

So, how do we do this? Let’s write a very simple custom dataset class first.

Also, let’s generate some random data which we will use with this custom Dataset.

Example of one random sequence and label. Each integer in the sequence corresponds to a word in the sentence.

We can use the custom dataset now using:

train_dataset = CustomTextDataset(X,y)

If we now try to use the Dataloader on this dataset with batch_size >1, we will get an error. Why is that?

train_dataloader = DataLoader(train_dataset,batch_size = 64, shuffle=False, num_workers=10)

for xb,yb in train_dataloader:

print(xb.size(),yb.size())

This happens because the sequences have different lengths, and our data loader expects our sequences of the same length. Remember that in the previous image example, we resized all images to size 224 using the transforms, so we didn’t face this error.

So, how do we iterate through this dataset so that each batch has sequences with the same length, but different batches may have different sequence lengths?

We can use collate_fn parameter in the DataLoader that lets us define how to stack sequences in a particular batch. To use this, we need to define a function that takes as input a batch and returns ( x_batch , y_batch ) with padded sequence lengths based on max_sequence_length in the batch. The functions I have used in the below function are simple NumPy operations. Also, the function is properly commented so you can understand what is happening.

We can now use this collate_fn with our Dataloader as:

train_dataloader = DataLoader(train_dataset,batch_size = 64, shuffle=False, num_workers=10,collate_fn = collate_text) for xb,yb in train_dataloader:

print(xb.size(),yb.size())

See that the batches have different sequence lengths now