ImageFolder dataLoader for ImageNet with selected classes and pretrained PyTorch model

username: jS5t3r
cooked: <p>I selected some classes by doing soft links into a <code>val</code> folder for the validation set.</p>
<pre><code class="lang-bash">class="n02841315"  # binoculars
ln -s /home/from/ImageNet/val/"$class" /home/to/ImageNetHierarchy/val/"$class"
</code></pre>
<p>I use this dataloader:</p>
<pre><code class="lang-python">import torch
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
transform = transforms.Compose(
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    normalize,
)

dataset_dir_path = "/home/to/ImageNetHierarchy/val/"
data_loader = torch.utils.data.DataLoader(datasets.ImageFolder(dataset_dir_path, transform), batch_size=64, shuffle=True, num_workers=num_workers, pin_memory=True)
</code></pre>
<p>When I try to evaluate a pretrained model such as</p>
<pre><code class="lang-auto">import torchvision
model = torchvision.models.wide_resnet50_2(pretrained=True)
model.eval()
model.cuda()

for images, labels in data_loader :
      images = images.cuda()
      labels = labels.cuda()

      outputs = model(images)
      _, predicted = torch.max(outputs.data, 1)
      
      if (predicted == labels):
          print("this hardly happens!")
</code></pre>
<p>This model hardly predicts the correct class.</p>

username: ptrblck
cooked: <p><code>ImageFolder</code> creates the class indices (i.e. the targets) based on the available folders.<br>
If I understand your use case correctly, you’ve only slimmed down the validation datasets while the training still has all 1000 classes?<br>
In that case you are corrupting the class correspondence as the mapping could look like (using random class names):</p>
<pre><code class="lang-python"># train
apple - class0
bird - class1
duck - class2
eagle - class3
wolf - class4

# val
bird - class0
eagle - class1
</code></pre>
<p>If you want to manipulate the validation dataset I would recommend to create a custom <code>Dataset</code> and make sure to provide the expected class labels to the remaining folders.</p>

username: jS5t3r
cooked: <p>yes, the training is still 1000 classes. The validation set has much less and that’s why it is corrupted.</p>
<p>What is the best way to create such a dataloader?</p>

username: Matias_Vasquez
cooked: <p>You can create your custom <code>Dataset</code> that returns the expected value corresponding to the original 1000 classes</p>
<pre data-code-wrap="python"><code class="lang-python">import torchvision

class MyImageFolder(torchvision.datasets.ImageFolder):
    def __init__(self, img_path, transform=None):
        super(MyImageFolder, self).__init__(img_path, transform)
        self.classes, self.class_to_idx = self._my_classes()
        self.samples = self._make_dataset(self.samples)
        self.imgs = self.samples
        self.targets = [s[1] for s in self.samples]

    def _my_classes(self):
        classes = ['duck', 'wolf']
        class_to_idx = {classes[i]: i for i in range(len(classes))}

        return classes, class_to_idx

    def _make_dataset(self, samples):
        n = len(samples)
        ds = [None] * n
        
        for i, (img, cls) in enumerate(samples):
            ds[i] = (img, self._custom_class(cls))

        return ds

    def _custom_class(self, cls):
        if cls == 0:
            return self.classes[0]
        if cls == 1:
            return self.classes[1]
        else:
            return 'not_my_favorite_class'
</code></pre>
<p>This would be a slight variation to the answer given here.</p>
<aside class="quote" data-post="5" data-topic="52300">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/custom-label-for-torchvision-imagefolder-class/52300/5">Custom label for torchvision ImageFolder class</a> <a class="badge-category__wrapper " href="/c/vision/5"><span data-category-id="5" data-drop-close="true" class="badge-category " title="Topics related to either pytorch/vision or vision research related topics"><span class="badge-category__name">vision</span></span></a>
  </div>
  <blockquote>
    The internal classes and class_to_idx attributes are used in DatasetFolder's __init__ to create the samples as seen in <a href="https://github.com/pytorch/vision/blob/d2c763e14efe57e4bf3ebf916ec243ce8ce3315c/torchvision/datasets/folder.py#L93-L94" rel="noopener nofollow ugc">these lines of code</a>. 
After ImageFolder was initialized (DatasetFolder is the parent class of it), changing these values won’t have any effect on the samples anymore, so you might want to derive your custom class from DatasetFolder and change these attributes in the __init__ method.
  </blockquote>
</aside>

<p>Hope this helps <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>