Writing fast Dataset class

username: Shalabh_Gupta
cooked: <p>Hi all! I am relatively new to pytorch, and wanted some help in writing a dataset class that could possibly use pytorch inbuilt dataset generators like DataFolder or other utilities.</p>
<p>I am working on the ModelNet10 dataset and my directory structure looks like follows:<br>
root directory contains 10 folders (bathtub,bedroom,toilet,…). These contain 2 folders (train and test) each, which further contain around 500 stl files. Right now , my dataset class looks like follows:</p>
<p>class MeshData(Dataset):</p>
<pre><code>def __init__(self,root_dir):
    self.classes = set()
    self.X = []
    self.Y = []
    self.final_dataset = []
    self.classes_codec = LabelEncoder()

    lst_of_classes = os.listdir(root_dir)
    lst_of_classes.sort()

    self.classes_codec.fit(lst_of_classes)

    for x in lst_of_classes:
        print(x)

        path = root_dir + '/' + x + '/train'
        lst_of_objects = os.listdir(path)
        for y in lst_of_objects:
            file_path = path + '/' + y

            dict_para = input_stl(file_path)
            if dict_para == False:
                print('Input file',y,'has problems')
                continue
            else:	
                neigh = dict_para["neigh_index"]
                corner = dict_para["corners"]
                center = dict_para["centroids"]
                normal = dict_para["normals"]

                self.X.append(np.concatenate((center,corner,normal,neigh),axis=1))
                self.Y.append((self.one_hot_encode(self.classes_codec,[x])))


    self.final_dataset = [self.X,self.Y]

def one_hot_encode(self,codec,values):
    value_idxs = codec.transform(values)
    val,idx = torch.max(torch.eye(len(codec.classes_))[value_idxs],1)
    return torch.LongTensor(idx)



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

def __getitem__(self,idx):
    return torch.from_numpy(self.X[idx]),self.Y[idx]
</code></pre>
<p>Here, i am running brute force over all folders and files, loading them in the format i want and one hot encoding the labels. But this is just too time consuming.<br>
I looked up some documentation and found torchvision.datasets.DatasetFolder to be doing a similar job, but on directory structure as follows:</p>
<p>root/class_x/xxx.ext<br>
root/class_x/xxy.ext<br>
root/class_x/xxz.ext</p>
<p>root/class_y/123.ext<br>
root/class_y/nsdf3.ext<br>
root/class_y/asd932_.ext</p>
<p>Only difference is I have train and test folders inside class_x and class_y and i only want to load training .stl files.</p>
<p>Thanks a lot.</p>

username: ptrblck
cooked: <p>The main difference between the <code>ImageFolder</code> and your approach would be the lazy vs. eager loading.<br>
While the <code>ImageFolder</code> dataset loads each sample lazily in the <code>__getitem__</code> method, it seems you are preloading all samples in the <code>__init__</code> method.</p>
<p>You could try to refactor your code such that you are only creating all valid paths (and targets, if it’s cheap to calculate them) in the <code>__init__</code> and execute the loading using these paths in <code>__getitem__</code>.</p>
<p>Let me know, if this would help.</p>

username: Shalabh_Gupta
cooked: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>, thanks for the reply. If I want to process the data efficiently in <strong>getitem</strong> method, i would have to keep the path and class name saved for each index idx in the <strong>init</strong> function itself, right? So that we could just lookup the idx easily.</p>
<p>Thanks.</p>

username: ptrblck
cooked: <p>Yes, that is correct.<br>
The usual approach would be to store the image paths and the targets in different lists and just index them in the <code>__getitem__</code> method with the provided <code>index</code> argument.</p>
<p>However, you would have to make sure that the lists containing the image path and target keep their correspondence. Sorting one of them would break your training.<br>
Alternatively you could also use a <code>dict</code>, and append each item as <code>{'img_path': path, 'target': target}</code>, but it depends on your coding style.</p>