Issue with custom Dataset class for DataLoader

username: ruchitpatel
cooked: <p>I’ve defined following class for my Dataset, so that I can later feed it to DataLoader class to generate my training data.</p>
<pre><code class="lang-auto">from torch.utils import data
import random
import numpy as np

class Dataset(data.Dataset):
    def __init__(self, qrel, query_vecs, offsets, fname):
        self.qrel = qrel
        self.query_vecs = query_vecs
        self.offsets = offsets
        self.f = open(fname, 'rt')
        self.keys = offsets.keys()
        
    def __len__(self):
        return len(self.qrel)
    
    def __iter__(self):
        return iter(self.qrel)
    
    def __next__(self):
        return next(iter(self.qrel))
    
    def __getitem__(self, key):
        # Following line fetches a query vector from the dictionary
        query = self.query_vecs[key]
        
        # Following four lines are responsible for fetching doc vector for positive document.
       # All doc vectors are stored as a tsv file
        posid = self.qrel[key][0]
        offset = self.offsets[posid]
        self.f.seek(offset)
        pos = np.fromstring(self.f.readline().split('\t')[1], sep=" ")
        
       #Following lines are responsible for fetching randomly selected four 
      # Document vectors
        negs = []
        neg_keys = random.sample(self.keys, 4)
        for key in neg_keys:
            self.f.seek(self.offsets[key])
            negs.append(np.fromstring(self.f.readline().split('\t')[1], sep=" "))
            
        # Return tuple of query, positive, negative docs.  
        return (query, pos, np.array(negs))
</code></pre>
<p>I’m using this class as follows :</p>
<pre><code class="lang-auto">import torch
from torch.utils import data
import pickle
import random
import numpy as np
import gzip
import csv
from my_dataset import Dataset

qrel = {}
with gzip.open("msmarco-doctrain-qrels.tsv.gz", 'rt', encoding='utf8') as f:
    tsvreader = csv.reader(f, delimiter="\t")
    for [topicid, _, docid, rel] in tsvreader:
        assert rel == "1"
        if topicid in qrel:
            qrel[topicid].append(docid)
        else:
            qrel[topicid] = [docid]
      
# query_vecs is a dictionary      
with open('query_vecs.pickle', 'rb') as f:
    query_vecs = pickle.load(f)

# Offsets is also a dictionary    
with open('doc_vec_offsets.p', 'rb') as f:
    offsets = pickle.load(f)
    
fname = 'doc_vecs.tsv'


params = {'batch_size': 1,
          'shuffle': False,
          'num_workers': 1}

max_epochs = 1
ds = Dataset(qrel,query_vecs, offsets, fname)
dataloader = data.DataLoader(ds, **params)

for epoch in range(max_epochs):
    for sample in dataloader:
        print(sample)
        print(type(sample))
        break
</code></pre>
<p>However, I’m getting following error. I don’t know why I’m getting this KeyError, because I’ve already replaced the <code>__iter__()</code> and <code>__next__()</code> method. What else should I do? Basically my dataset starts from that qrel dictionary, where it first finds the query vector, positive document vector and 4 negative doc vectors. However, I can not just iterate through that qrel in the class Dataset. I get following error and I’m stuck.</p>
<pre><code class="lang-auto">
Exception: KeyError:Traceback (most recent call last):
  File "/home/ruchit/anaconda3/lib/python3.7/site-packages/torch/utils/data/_utils/worker.py", line 99, in _worker_loop
    samples = collate_fn([dataset[i] for i in batch_indices])
  File "/home/ruchit/anaconda3/lib/python3.7/site-packages/torch/utils/data/_utils/worker.py", line 99, in &lt;listcomp&gt;
    samples = collate_fn([dataset[i] for i in batch_indices])
  File "/home/ruchit/Desktop/NLP/project/my_dataset.py", line 24, in __getitem__
    query = self.query_vecs[self.qrel[key]]
KeyError: 0
</code></pre>
<p>It looks like even after I’ve defined <strong>iter</strong>() method, it still is replacing it with iterator corresponding to some list.<br>
The problem seems to be with DataLoader class. I can iterate through object of class <code>Dataset</code>, but after doing</p>
<pre><code class="lang-auto">ds = Dataset(qrel,query_vecs, offsets, fname)
dataloader = data.DataLoader(ds, **params)
</code></pre>
<p>dataloader is no longer iterable.</p>
<p>For example,<br>
below code works fine</p>
<pre><code class="lang-auto">for sample in ds:
    print(sample)
</code></pre>
<p>But, after passing it to DataLoader, it not longer works. Following thing fails :</p>
<pre><code class="lang-auto">for sample in dataloader:
    print(sample)
</code></pre>

username: Stacie_Davis
cooked: <p>Check if <code>next(iter(x))</code> inside your <code>__next__</code> is causing problem , because essentially for each call to <code>next</code>, you are re-initializing iterator.</p>

username: ruchitpatel
cooked: <p>At first, I did not implement the <strong>next</strong> method at all, but I was getting the same error. I thought implementing it should remove it. But should I remove the next method at all from my code?</p>

username: ruchitpatel
cooked: <p>I did one hack here. Instead of iterating over dictionary, I iterated over its keys. Then I don’t have to implement the <strong>iter</strong>() pattern any more. Its working now, but there is a IndexError when I run it for 8 workers. For one worker, it works fine.</p>

username: dhpollack
cooked: <p>Have you tried it without the <code>__iter__</code> and <code>__next__</code> methods?  Generally, you are only supposed to only override <code>__getitem__</code> for datasets as in the <a href="https://pytorch.org/tutorials/beginner/data_loading_tutorial.html" rel="nofollow noopener">tutorial</a>.</p>

username: ruchitpatel
cooked: <p>Yes, as I mentioned in previous comment, I just iterated over dictionary keys, where I removed the iter and next methods. I managed to solve this problem. But after this, I am not able to set num_workers=8, it only works for num_workers=1. When I set it to 8, I get the IndexError.  See <a href="https://discuss.pytorch.org/t/indexerror-when-running-dataloader-for-multiple-workers/80401">this</a> for details. I guess I should close this ticket if it’s possible. I’ll check if I can do that.</p>