How to represent class_to_idx map for custom dataset in Pytorch

I am attempting transfer learning with a CNN (vgg19) on the Oxford102 category dataset consisting of 8189 samples of flowers labeled from 1 through 102. Instead of loading the data with ImageFolder, which requires a tedious process of structuring my data into train, valid and test folders with each class being a sub-folder holding my images, I decided to load it in using the Custom Dataset class following This tutorial

My directory paths set up for the image, the labels (from 1 through 102), and class label map which
maps the label to flower name

    data_dir_path = 'data/images/'
    labels_path = 'data/imagelabels.mat'
    class_label_path = 'data/class_label_map'y 

my data is not in the ImageFolder, and as far as I know If writing a custom dataset, you don’t need your data to be in that format.


my labels list which is in data/imagelabels.mat is as follows,

[77 77 77 … 62 62 62]

The transforms I am using for train, test and valid are

    data_transforms = {
'train': transforms.Compose([
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225])
'valid': transforms.Compose([
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225])
'test': transforms.Compose([
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225])


Here is my code to load in data, it works as I can, index into my images and labels.

    class MyDataset(Dataset):

def __init__(self, image_labels, data_dir, transform=None):


    :param image_labels_path: path to our labels
    :param root_dir: the directory which houses our images
    :param transform: apply any transform on our sample

    self.image_labels = image_labels
    self.root_dir = data_dir
    self.transform = transform

def __len__(self):

    label_dict =
    return len(label_dict['labels'][0]) # [0] cause its a list of lists

def __getitem__(self, idx):

    image_path_list = [os.path.join(self.root_dir, filename) for filename in os.listdir(self.root_dir)]

    image =[idx])

    label_dict =
    label_list = label_dict['labels'][0]

    # labels for pytorch should start from zero
    label_list[:] = [i - 1 for i in label_list]
    label = label_list[idx]  

    if self.transform:
        image = self.transform(image)

    return image, label

Here I am creating my trainset, validset and testset

    image_datasets = {x: MyDataset(image_labels=labels_path, data_dir=data_dir_path, transform=data_transforms[x]) for x in
              ['train', 'valid', 'test']}

I subtracted 1 from my label list since Pytorch expects labels to start from 0. Thus 0 through 101 for 102 labels.

Correct me if I am wrong cause I get a “current target >=0 and current target <= n_classes failed” error if I don’t subtract one. Since this fails at 102, I thus surmised that class representation starts from 0.

I have a class_label_map is a dict which maps class labels to flowers names

{ 1": “pink primrose”,
“2”: “hard-leaved pocket orchid”,
“3”: “canterbury bells”,
“4”: “sweet pea”,
“5”: “english marigold”,
“6”: "tiger lily }

My big problem is to get a class_to_idx mapping, how do I do this, my flower names do not match the images if I visualize them, I get totally different flower names for my flowers.

I first created a mapping by having a dict with key my original label before I subtracted 1, and value the one after. Example:

        class_to_idx = {77:76, 73:72, 1:0, 65:64......102:101...65:54}

This is beyond doubt wrong as I was getting totally wrong label for my images.

The first label of my image in my data_dir_path = ‘data/images’ is 77, upon subtracting one I get 76. Would this mean the index for all labels 76 is 0, and if the next class is 72 , would that mean the index for all classes 72 is 1? So…

        class_to_idx = {76:0, 72:1, 0:2, 65:3....and so on}

But this again works if the class number does not repeat after a certain classes, i.e. if 72 were to repeat
after class 65, then 72 is no more the 2nd class. The code for how ImageFolder find class_to_index is…

     classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]
     class_to_idx = {classes[i]: i for i in range(len(classes))}

I don’t properly understand this. Performing it on my label_list give me the largest index of the a certain
class {1: 39, 2: 99, 3: 139, 4: 195, 5: 260, 6: 305, 7: 345, 8: 430, 9: 476}.

The ImageFolder seems to have a class_to_idx attribute which if used on my Dataset throws an error,


AttributeError: ‘MyDataset’ object has no attribute ‘class_to_idx’

This is obviously the case because my Dataset class does not contain any such attribute. If I am reasoing is correct, Pytorch gives out a tuple of probabilities and corresponding index when we do probs.max(dim=1). This index is the value in our class_to_index map. The key is the also the key in our
label_map with, the value of which helps us trace the flower name of the our prediction.

So How do I map my classes to my index in this scenario? It isn’t structured like ImageFolder.
This is super important as I need to checkpoint my model and load it back again to fire out predictions.It may sound really silly, but I really don’t know what to do here, please please please help?

Yes, you are correct in your assumption that Python uses 0-based indices.
Your labels should indeed be in the range [0, nb_classes-1].

Since your map starts at the index 1, you should just add 1 to your label:

idx_to_class = {
    1: "pink primrose",
    2: "hard-leaved pocket orchid",
    3: "canterbury bells",
    4: "sweet pea",
    5: "english marigold",
    6: "tiger lily"

label = 0
sample_class = idx_to_class[label+1]

No need to create another mapping.

Also as a small side note: it seems you are using the whole data for your training, validation, and test datasets, since you are just changing the transformations.

Thanks a lot, the thing I am trying to grok here is that Pytorch doesn’t actually return the label of the image but the index of the highest probability in the output tensor right? Does it internally sort classes from 0 through n-1 classes such that index correspond to the class.

So though my label_list doesn’t start with class 0 but class 77 , and suppose my model returns index 48 as the highest probability when I do probs.max(dim=1)[1] by passing in some random image, does the index 48 mean class 48, which I can then look in my class_dict . Hope my question is clear and thank you once again.

Regarding the whole data for train, valid and test, I have a function that does the splitting into each.

    def get_train_valid_split(dataset, shuffle=True, first_split=0.8, second_split=0.9, random_seed=1):


:param dataset: dataset
:param shuffle: shuffles the indices so random elements from our data find themselves in our splits.
:param first_split: this is for train
:param second_split: for validation
:param random_seed: for test
:return: train, validation and test loaders
dataset_size = len(dataset)  # get any, we just need the length, say its n
indices = list(range(dataset_size))  # 0 though n-1.
split1 = int(np.floor(first_split * dataset_size))
split2 = int(np.floor(second_split * dataset_size))

if shuffle:

train_idx = indices[:split1]
val_idx = indices[split1:split2]
test_idx = indices[split2:]

samplers = {'train': SubsetRandomSampler(train_idx), 'valid': SubsetRandomSampler(val_idx),
            'test': SubsetRandomSampler(test_idx)}

dataloader = {x:[x], batch_size=64, sampler=samplers[x]) for x in
               ['train', 'valid', 'test']}

return dataloader

data_loaders = get_train_valid_split(image_datasets[‘train’]) # any key works, need to get length in order to split.

   {'train': < object at 0x7f7966ceb860>, 'valid': < object at 0x7f7966ceb978>, 'test': < object at 0x7f7966ceb9e8>}

    dataset_sizes = {x: len(data_loaders[x]) for x in ['train', 'valid', 'test']}
   {'train': 103, 'valid': 13, 'test': 13} # batch_size = 64.

probs.max(dim=1)[1] will give you the predicted class index. This class index does not inherit any class information, but will be learned by providing the right labels.
The ground truth label is provided by the Dataset.
Could you explain a bit more about imagelabels.mat?
Is it just a MATLAB matrix containing indices?
If so, do you know how the mapping was created?
Currently it looks like you are trying to map the sorted file paths to the values in imagelabels.mat, which might be wrong. E.g. does image_00001.jpg belong to class77 (MATLAB index)?

Thanks for the info regarding the datasets. I just thought this could be another issue after fixing this one here :wink:

The imagelabels.mat is an array of labels.

    label_list = [77 77 77 ... 62 62 62]

    <class 'numpy.ndarray'>

The sorted version of my file path is as follows

    image_path_list = sorted([os.path.join('data/images/', filename) for filename in os.listdir('data/images/')])

This gives
[‘data/images/image_00001.jpg’, ‘data/images/image_00002.jpg’, ‘data/images/image_00003.jpg’, ‘data/images/image_00004.jpg’…and so on]

The corresponding class indeed is 77. And it maps perfectly to Passion flower. A few other randomly tested map too.
Me being the very embodiment of stupidity did not sort my labels. And the order they were in since I not sorted them was totally arbitrary as os.listdir() doesn’t read files in the order they appear in the directory.

So, I ended up subtracting one from label_list with 77 becoming 76 and 73 72 and so on. Create a
new class_to_label map that starts from 0.

"0": "pink primrose",
"1": "hard-leaved pocket orchid",
"2": "canterbury bells",
"3": "sweet pea",
"4": "english marigold",
"5": "tiger lily",
"6": "moon orchid",
"7": "bird of paradise",
"8": "monkshood",
"9": "globe thistle",
"10": "snapdragon",
"11": "colt's foot",
"12": "king protea",
 and so on...
 "76": "passion flower"


I now think I am good to go. Thanks a ton.

Awesome it’s working!
So the issue was that your file paths weren’t sorted?

Yes that was the issue. I ended up printing out a few lines to check. Should have done that before. Thanks.

But going forward, are you aware of any best practice with this class_to_idx thing for Pytorch. This was the one that tripped me up the most. Thank you very much for your time and patience.

I have an issue using CIFER and rotation. I have a secondary information table I load which I wish to cross reference with the image I get from torch.index_select. How on earth do I get the Tensor index (as in it’s position in the dataset from say 0 - 50000)??

You can write a custom Dataset (e.g. by deriving from the CIFAR dataset) as described e.g. here and return the index in the __getitem__ method with the data and target tensors.

Ok, this may sound stupid but how do I use the get item (easy to change as I have modified the CIFAR loader to work with my data) just not sure the code to go from index_select to sub data
(Thanks for the reply)

Something like this would work:

class MyCIFAR10(
    def __init__(self, transform=None):
        self.dataset = datasets.CIFAR10(root='~/python/data', download=False, transform=transform)
    def __len__(self):
        return len(self.dataset)
    def __getitem__(self, index):
        x, y = self.dataset[index]
        return x, y, index
dataset = MyCIFAR10(transform=transforms.ToTensor())
loader =, batch_size=2)
x, y, index = next(iter(loader))
# > tensor([0, 1])