Speed Issues with shuffling in a torch IterableDataset

Viewed 67

On Google Colab I built an IterableDataset for loading the approx. 200K pics of the CELEBA dataset, which are located in form of 4 zip files on my GoogleDrive. I process them further in a GAN training loop. The code performs 1.3 iterations/sec on the GPU as long as I don't shuffle the filenames of the zips, but it gets incredibly slow, when I shuffle them shuffle=True (more than 10 sec. for each iteration). It is not the shuffling itself fnames = [fnames[i] for i in torch.randperm(len(fnames))] , which I tried out seperately and it is equally fast as without shuffling, while I run an empty training loop, it only slows down together with the full GAN training.

class CELEBA_Dataset_multi(IterableDataset):
  def __init__(self, zipnames, PATH="/content/drive/MyDrive/LargeFiles/celeb_dataset/", shuffle=False):
    self.zipnames = [PATH + zipname for zipname in zipnames]
    self.shuffle = shuffle
    self.len = 0    
    for zipname in self.zipnames:
      with ZipFile(zipname, 'r') as ZipImg:
        self.len += len(ZipImg.namelist())
  def parse_zipfiles(self, zipnames):
    trf = transforms.Compose(
        [transforms.Resize(64),
         transforms.CenterCrop(64),
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])    
    for zipname in self.zipnames:
      with ZipFile(zipname, 'r') as ZipImg:
        fnames = ZipImg.namelist()
        if self.shuffle: fnames = [fnames[i] for i in torch.randperm(len(fnames))]
        for name in fnames:
          img = ZipImg.open(name)
          yield trf(Image.open(img))
  def __iter__(self):    
      return self.parse_zipfiles(self.zipnames)
  def __len__(self):
    return self.len
    

Does anybody have an idea, what is wrong?

0 Answers
Related