Hey there! It makes total sense that your lighter model shifted the bottleneck to data loading—when the model can process batches faster than your disk can feed them, the GPU ends up sitting idle. Let’s walk through actionable fixes tailored exactly to your setup:

1. Pre-Crop & Cache Smaller Images Offline #

Your current workflow loads 1040x1920 images every time just to crop them down to 256x256, which is a massive waste of IO and CPU time. Instead, pre-generate cropped versions of your dataset once:

• Write a quick script to iterate over all real/fake images, apply RandomCrop(256) 5-10 times per original image (to preserve randomness), and save the cropped 256x256 PNGs to new directories like real_cropped/ and fake_cropped/.
• Update your TrainImageDataset to load these pre-cropped images directly. You can keep the random horizontal flip in the dataset if needed—flipping is trivial compared to loading/cropping large images.

This cuts your disk read size by ~90% (from ~6MB per full-size RGB image to ~0.2MB per cropped one), drastically speeding up data loading.

2. Optimize DataLoader Parameters #

Tweak these settings to get more out of your existing pipeline:

• Increase num_workers: Set it to match or double your CPU core count (e.g., 8 workers for an 8-core CPU). More workers mean parallel disk reads and preprocessing.
• Enable persistent_workers=True: By default, DataLoader destroys and recreates worker processes between epochs—this adds unnecessary overhead. Keeping workers persistent lets them stay alive across epochs, saving time on process startup.
• Adjust prefetch_factor: Set this to 2 or higher (e.g., prefetch_factor=4) to make workers preload the next 2-4 batches while the GPU processes the current one.

Updated DataLoader example:

dataloader = DataLoader(
    dataset,
    batch_size=8,
    num_workers=8,  # Match your CPU core count
    shuffle=True,
    pin_memory=True,
    drop_last=True,
    persistent_workers=True,
    prefetch_factor=4
)

3. Use Faster Image Loading Libraries #

Replace PIL with faster alternatives to speed up image decoding:

• Pillow-SIMD: A drop-in replacement for PIL with SIMD optimizations that can speed up image loading by 2-4x. Install it with pip install pillow-simd—your existing code will work without changes since the API is identical.
• TorchVision IO: Use torchvision.io.read_image instead of PIL to load images directly as tensors, skipping the F.to_tensor step. Modify your __getitem__ like this:

  from torchvision.io import read_image
  
  def __getitem__(self, batch_index):
      # Load directly as tensor and normalize to [0,1]
      real = read_image(self.real_images[batch_index]) / 255.0
      fake = read_image(self.fake_images[batch_index]) / 255.0
      # Apply augmentations
      fake = self.downscale(fake)
      real = self.downscale(real)
      fake = self.hflip(fake)
      real = self.hflip(real)
      return {'fake': fake, 'real': real}
  

4. Cache Data in Memory (For Smaller Datasets) #

If your total dataset size (after cropping) fits in system memory, load all images into memory during dataset initialization:

class TrainImageDataset(Dataset):
    def __init__(self, path_real, path_fake, img_size=256):
        super().__init__()
        self.downscale = transforms.RandomCrop(img_size)
        self.hflip = transforms.RandomHorizontalFlip(p=0.5)
        
        # Load all images into memory upfront
        self.real_images = []
        for img_path in [os.path.join(path_real, x) for x in os.listdir(path_real)]:
            self.real_images.append(Image.open(img_path))
        
        self.fake_images = []
        for img_path in [os.path.join(path_fake, x) for x in os.listdir(path_fake)]:
            self.fake_images.append(Image.open(img_path))

    def __getitem__(self, batch_index):
        real = self.real_images[batch_index]
        fake = self.fake_images[batch_index]
        # Apply augmentations
        fake = self.downscale(fake)
        real = self.downscale(real)
        fake = self.hflip(fake)
        real = self.hflip(real)
        real = F.to_tensor(real)
        fake = F.to_tensor(fake)
        return {'fake': fake, 'real': real}

This eliminates disk reads entirely during training—perfect if you have enough RAM (e.g., 10k cropped images take ~2GB).

5. Try NVIDIA DALI (For Advanced GPU-Accelerated Loading) #

If you want to push performance further, NVIDIA DALI lets you offload preprocessing (including cropping large images) directly to the GPU without loading full-size images into CPU memory. It integrates seamlessly with PyTorch and handles the entire data pipeline from disk to GPU in parallel. While there’s a small learning curve, it’s worth it for large datasets with heavy preprocessing.

Recommended Order of Testing #

Start with the easiest, lowest-effort fixes first:

1. Install Pillow-SIMD and update DataLoader parameters (persistent_workers, num_workers)
2. Pre-crop your dataset offline
3. If still not enough, try in-memory caching or DALI

内容的提问来源于stack exchange，提问作者arturo salmi