Dataloader iterates through entire dataset

DrJellybean: <p>I created a dataset that loads a single data sample at a time on demand (1 sample consists of multiple images), and I have a data loader with a small batch size. When I try to show just the first few batches of my dataset, the loader keeps trying to iterate through my entire dataset instead of pulling out just the few data samples:</p>
<pre><code class="lang-auto">class FaceDataset(Dataset):
    def __init__(self):
        df = pd.read_csv("data/positions.csv")
        df["filename"] = df["id"].astype("str") + ".jpg"

        self.filenames = df["filename"].tolist()
        self.targets = torch.FloatTensor(list(zip(df["x"], df["y"])))
        self.head_angle = torch.FloatTensor(df['head_angle'].tolist())

    def __len__(self):
        return len(self.targets)

    def __getitem__(self, idx):
        sample = {
            "targets": self.targets[idx],
            "head_angle": self.head_angle[idx]
        }

        for img_type in ["face", "face_aligned", "l_eye", "r_eye", "head_pos"]:
            img = Image.open("data/{}/{}".format(img_type, self.filenames[idx]))
            img = torch.from_numpy(np.array(img))
            # img /= 255
            sample[img_type] = img

        return sample

ds = FaceDataset()
data = DataLoader(ds, batch_size=2, shuffle=True, num_workers=2)

for i_batch, sample_batched in enumerate(data):
    print(i_batch, sample_batched)

    if i_batch == 1:
        break
</code></pre>
<p>The length of my dataset is about 15k, and the loader seems to be trying to load everything instead of just a single batch of 2. Jupyter seems to just freeze before anything actually gets printed. Am I creating my dataset/loader incorrectly?</p>
<p>When I try to iterate over the dataset, I seem to get the correct data back for 1 training sample. The problem only seems to come in when iterating over the data loader object:</p>
<pre><code class="lang-auto">for batch in ds:
    print(batch)
    break
</code></pre>
<p>EDIT: I think Ive narrowed this down to the <code>num_workers</code> param. When I set it to <code>0</code> then the whole thing works as expected, but when I set it to <code>&gt; 0</code> then nothing gets printed?</p>

ptrblck: <p>Your small code snippet runs correctly on my machine:</p>
<pre><code class="lang-python">dataset = TensorDataset(torch.randn(100, 1), torch.randn(100, 1))
loader = DataLoader(dataset, batch_size=2, shuffle=True, num_workers=2)

for batch_idx, batch in enumerate(loader):
    print(batch_idx, batch)
    if batch_idx == 1:
        break

&gt; 0 [tensor([[ 0.7712],
          [-0.3303]]), tensor([[0.0505],
          [1.3738]])]
  1 [tensor([[-0.5104],
          [-0.5665]]), tensor([[-0.4806],
          [-0.7598]])]
</code></pre>
<p>Is the <code>DataLoader</code> with multiple workers working at all in your current setup and is only this small loop creating issues?</p>

dashesy: <p>I see same error now, even with <code>num_workers</code> of 0 and using <code>DistributedSampler</code> when using it on a <code>world_size</code> of 1</p>

ptrblck: <p>Could you post a minimal and executable code snippet reproducing the issue?</p>

dashesy: <p>It was my fault! Custom dataset was throwing an error I did not catch!</p>