Iterate with DataLoader over multiple datasets with shuffle enabled

Lelelon: <p>I am using the huggingface library and PyTorch, hopefully this question is best suited here.</p>
<p>I have tokenized a dataset in two different ways resulting in two different tokenized datasets. I want to iterate over the datasets with shuffling enabled, such that the batches correspond to the same examples. I created a small example of the problem I am facing:</p>
<pre><code class="lang-python">import torch
from transformers import AutoTokenizer
from datasets.load import load_dataset
from torch.utils.data import DataLoader
from transformers import DataCollatorWithPadding

data = load_dataset("glue", "mnli", split='train[:4]')

checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

# Tokenize data
tokenized_data = data.map(lambda example: tokenizer(
            example["premise"], example["hypothesis"], truncation=True))
tokenized_data = tokenized_data.remove_columns(["premise", "hypothesis"])
tokenized_data = tokenized_data.rename_column("label", "labels")
tokenized_data.set_format("torch")

h_tokenized_data = data.map(lambda example: tokenizer(
            example["hypothesis"], truncation=True))
h_tokenized_data = h_tokenized_data.remove_columns(["premise", "hypothesis"])
h_tokenized_data = h_tokenized_data.rename_column("label", "labels")
h_tokenized_data.set_format("torch")

# Create loaders
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
train_loader = DataLoader(tokenized_data, batch_size=2, shuffle=True, collate_fn=data_collator)
h_train_loader = DataLoader(h_tokenized_data, batch_size=2, shuffle=True, collate_fn=data_collator)

for batch in zip(train_loader, h_train_loader):
    break

# I want the output to be the same:
print(batch[0]["idx"]) # out: tensor([2, 1])
print(batch[1]["idx"]) # out: tensor([3, 1])
</code></pre>
<p>I guess the most interesting solution would allow for different tokenizers and collators used, but anything that can help the particular example case I provided is appreciated.</p>

rinkujadhav2013: <p>I’m not sure if I understood everything but if you create 2 dataloaders with shuffle=True, they’re both gonna get shuffled randomly and differently.<br>
What you could try is set the random seeds to a fixed same value for both dataloaders using <code>worker_init_fn</code> as described <a href="https://pytorch.org/docs/stable/data.html#randomness-in-multi-process-data-loading" rel="noopener nofollow ugc">here</a>.<br>
Also check this whole thread: <a href="https://discuss.pytorch.org/t/fix-seed-for-data-loader/146711" class="inline-onebox">Fix seed for data loader</a></p>

Lelelon: <p>Specifying identical generators for the two dataloaders seems to do the trick:</p>
<pre><code class="lang-python">import torch
from transformers import AutoTokenizer
from datasets.load import load_dataset
from torch.utils.data import DataLoader
from transformers import DataCollatorWithPadding

from torch.utils.data.sampler import RandomSampler

data = load_dataset("glue", "mnli", split='train[:4]')

checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

# Tokenize data
tokenized_data = data.map(lambda example: tokenizer(
            example["premise"], example["hypothesis"], truncation=True))
tokenized_data = tokenized_data.remove_columns(["premise", "hypothesis"])
tokenized_data = tokenized_data.rename_column("label", "labels")
tokenized_data.set_format("torch")

h_tokenized_data = data.map(lambda example: tokenizer(
            example["hypothesis"], truncation=True))
h_tokenized_data = h_tokenized_data.remove_columns(["premise", "hypothesis"])
h_tokenized_data = h_tokenized_data.rename_column("label", "labels")
h_tokenized_data.set_format("torch")

# Added: Create generators with the same seed to load the data in the same order
SEED = 1
generator1 = torch.Generator()
generator1.manual_seed(SEED)
generator2 = torch.Generator()
generator2.manual_seed(SEED)

# Create loaders
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
train_loader = DataLoader(tokenized_data, batch_size=2, collate_fn=data_collator, shuffle=True, generator=generator1)
h_train_loader = DataLoader(h_tokenized_data, batch_size=2, collate_fn=data_collator, shuffle=True, generator=generator2)

for batch in zip(train_loader, h_train_loader):
    break

print(batch[0]["idx"]) # out: tensor([0, 3])
print(batch[1]["idx"]) # out: tensor([0, 3])
</code></pre>