Custom dataset not working on colab but works locally

username: Norrlandsfarger
cooked: <p>Hi, I was loading a custom dataset of images using Google Colab, but somehow it’s unable to recognise each element in the dataset properly. The code works fine locally. In Colab the variable list in colab does not show the shapes of <code>food_test</code> and <code>food_train</code> (should be 1000 and 9000):</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/b/bb4005d8402ba2c3a968a870dd9a2ce16cb82d6c.png" data-download-href="https://discuss.pytorch.org/uploads/default/bb4005d8402ba2c3a968a870dd9a2ce16cb82d6c" title="Screenshot from 2024-07-08 23-47-10"><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/b/bb4005d8402ba2c3a968a870dd9a2ce16cb82d6c.png" alt="Screenshot from 2024-07-08 23-47-10" data-base62-sha1="qIuse0qg1XCLqgzGoRxSTMgPPek" width="690" height="156" data-dominant-color="4D4D4D"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot from 2024-07-08 23-47-10</span><span class="informations">947×215 20.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>But the same code runs without issue locally.</p>
<p>I tried to <code>print(len(food_train))</code>, and it does print the right value, but if I attempt to check each one in the dataset with  <code>print(next(iter(train_loader)))</code>, it would go into an infinite loop.</p>
<p>This is my custom dataset with related functions:</p>
<pre><code class="lang-auto">def food_colormap2label():
    food_colormap = [(i, 0, 0) for i in range(103)]
    colormap2label = torch.zeros(256 ** 3, dtype=torch.long)
    for i, colormap in enumerate(food_colormap):
        colormap2label[
            (colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i
    return colormap2label

def food_label_indices(colormap, colormap2label):
    colormap = colormap.permute(1, 2, 0).numpy().astype('int32')
    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256
           + colormap[:, :, 2])
    return colormap2label[idx]

def food_rand_crop(feature, label, height, width):
    rect = torchvision.transforms.RandomCrop.get_params(
        feature, (height, width))
    feature = torchvision.transforms.functional.crop(feature, *rect)
    label = torchvision.transforms.functional.crop(label, *rect)
    return feature, label

class FoodSegDataset(torch.utils.data.Dataset):
    def __init__(self, root, is_train, crop_size):
        self.transform = torchvision.transforms.Normalize(
            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        self.crop_size = crop_size
        self.image_paths, self.label_paths = read_image_paths(root, is_train=is_train)
        self.colormap2label = food_colormap2label()

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

    def __getitem__(self, idx):
        image = torchvision.io.read_image(self.image_paths[idx])
        label = torchvision.io.read_image(self.label_paths[idx])
        # Check if the image is large enough for the crop size
        if image.shape[1] &gt;= self.crop_size[0] and image.shape[2] &gt;= self.crop_size[1]:
            image, label = food_rand_crop(image, label, *self.crop_size)
            image = self.transform(image.float() / 255.)
            label = food_label_indices(label, self.colormap2label)
            return image, label
        else:
            # If the image is too small, skip this example by returning a None tuple
            return self.__getitem__((idx + 1) % len(self))
</code></pre>
<p>And this is the loading part:</p>
<pre><code class="lang-auto">curr_dir = '/content/drive/MyDrive/data'

food_classes = read_food_categories(os.path.join(curr_dir, 'category.txt'))
food_colormap = [(i, 0, 0) for i in range(103)]

crop_size = (320, 480)
food_train = FoodSegDataset(curr_dir, True, crop_size)

food_test = FoodSegDataset(curr_dir, False, crop_size)

batch_size = 256
train_loader = torch.utils.data.DataLoader(food_train, batch_size, shuffle=True,drop_last=True, num_workers=1)
test_loader = torch.utils.data.DataLoader(food_test, batch_size, shuffle=False, drop_last=True, num_workers=1)

print(len(food_train))
# print(next(iter(train_loader)))
</code></pre>
<p>Thanks in advance!</p>

username: Brock_Brown
cooked: <p>What are you looking for if <code>len(food_test)</code> and <code>len(food_train)</code> are working? You’re wanting to see the value in the Colab debugger instead? Also regarding your infinite loop (I suspect you’re probably just seeing the first element over and over again), it might have something to do with your use of <code>iter</code> in <code>next(iter(train_loader)</code>:</p>
<pre data-code-wrap="python"><code class="lang-python"># ok usage of iter; create an iterable then use next to grab the next element
foo = [1, 2, 3]
bar = iter(foo)
for i in range(len(foo)):
    print(next(bar))
# output:
# 1
# 2
# 3
</code></pre>
<pre data-code-wrap="python"><code class="lang-python"># bad usage of iter
foo = [1, 2, 3]
bar = iter(foo)
for i in range(len(foo)):
    # instead of getting the next thing from the iterable, this will
    # create a new iterable then grab the first element
    print(iter(next(bar)))
# output:
# 1
# 1
# 1
</code></pre>

username: Norrlandsfarger
cooked: <p>Hi, thank you for the reply. For the <code>print(next(iter(train_loader)))</code>, the infinite loop is actually not showing any output. I tried to plot the images but the same happened. It says where the program got stuck but I’m not sure why: The exact same code is working fine on my local machine.</p>
<p>And if I ignore this then go straight to training, it would also get stuck without any output (loss for current epoch, running time, etc,.)</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/c/bce703e8118daacc6043aa1788502ff73ee42b37.png" data-download-href="https://discuss.pytorch.org/uploads/default/bce703e8118daacc6043aa1788502ff73ee42b37" title="Screenshot from 2024-07-09 08-22-44"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/c/bce703e8118daacc6043aa1788502ff73ee42b37_2_558x500.png" alt="Screenshot from 2024-07-09 08-22-44" data-base62-sha1="qX6IaimMeuKDdsVa49TS4Gyv8tF" width="558" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/c/bce703e8118daacc6043aa1788502ff73ee42b37_2_558x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/c/bce703e8118daacc6043aa1788502ff73ee42b37_2_837x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/c/bce703e8118daacc6043aa1788502ff73ee42b37_2_1116x1000.png 2x" data-dominant-color="313030"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot from 2024-07-09 08-22-44</span><span class="informations">1529×1368 238 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: Norrlandsfarger
cooked: <p><code>print(len(food_train))</code> works fine on colab, but it can’t print each element in the dataset:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/7/f77296f52c5b853074afda7b50b3a01056ffae35.png" data-download-href="https://discuss.pytorch.org/uploads/default/f77296f52c5b853074afda7b50b3a01056ffae35" title="Screenshot from 2024-07-08 23-41-03"><img src="https://discuss.pytorch.org/uploads/default/original/3X/f/7/f77296f52c5b853074afda7b50b3a01056ffae35.png" alt="Screenshot from 2024-07-08 23-41-03" data-base62-sha1="zj1rqNoNJrIsibRBKI4hlmJK7fn" width="690" height="199" data-dominant-color="2F302F"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot from 2024-07-08 23-41-03</span><span class="informations">747×216 12.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>