FSDP takes a very long time when start iterate DataLoader of each epoch

sxlllslgh: <p>I’m trying to train a CLIP-based model using FSDP on three GPUs, but when each epoch start iterating the data loader, it may take about 1 minute for reading 4K pieces of data.<br>
DataSet key  codes:</p>
<pre><code class="lang-auto">class MyDataset(DataSet):
    def __init__(self, …):
        self.label_table = pd.read_csv(…) # About 4K rows

    def __len__(self):
        return self.y.shape[0]
    
    def __getitem__(self, idx):
        img = Image.open(join(self.img_dir, self.label_table.loc[idx, 'image']))
        if self.transform is not None:
            img = self.transform(img)
        return img, …
</code></pre>
<p>DataLoader related codes:</p>
<pre><code class="lang-auto">transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((image_size, ) * 2),
    torchvision.transforms.ToTensor()
])

train_ds = MyDataset(…, transform=transform)
train_sampler = torch.utils.data.distributed.DistributedSampler(train_ds, rank=rank, num_replicas=args.world_size)
train_loader = torch.utils.data.DataLoader(train_ds, batch_size=args.batch_size, pin_memory=False, num_workers=48, sampler=train_sampler, drop_last=True)
</code></pre>
<p>Trainning codes:</p>
<pre><code class="lang-auto">model = MyModel(train_ds.get_ti_rads_size()).float().to(rank)
model = FullyShardedDataParallel(model, auto_wrap_policy=my_auto_wrap_policy, cpu_offload=CPUOffload(offload_params=True), use_orig_params=True)

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

for epoch in range(1, args.epochs + 1):
    fsdp_loss = torch.zeros(2).to(rank)
    if train_sampler:
        train_sampler.set_epoch(epoch)

    if rank == 0:
        epoch_progress = tqdm.tqdm(range(len(train_loader)), colour='blue', desc=f'Training Epoch {epoch}')
    # The hanging out problem appears on the line below of each epoch
    for image, …, _ in train_loader:
        # Copy tensors to devices.
        image = image.to(rank)
        # …
</code></pre>
<p>The program will stuck at the <code>for image, …, _ in train_loader:</code> for about 1 minute of each epoch.<br>
When I add debug print line at the <code>__getitem__</code> function of the DataSet, it seems that all images will be read again when DataLoader is iterated every time. So I would like to know how too avoid this problem?</p>
<p>I have tried some approaches but no one worked:</p>
<ol>
<li>Set <code>pin_memory=True</code>.</li>
<li>Read all images in the constructor function of DataSet, and read image tenors directly in the <code>__getitem__</code> function. (It might illustrate the DataSet would be reload at every epoch.</li>
</ol>

ptrblck: <aside class="quote no-group" data-username="sxlllslgh" data-post="1" data-topic="192084">
<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/sxlllslgh/48/28057_2.png" class="avatar"> sxlllslgh:</div>
<blockquote>
<p>but when each epoch start iterating the data loader, it may take about 1 minute for reading 4K pieces of data.</p>
</blockquote>
</aside>
<p>Your current <code>MyDataset.__init__</code> method only loads <code>self.label_table</code> via <code>pd.read_csv</code>, so I assume you are referring to this line of code?<br>
This behavior is expected since each new <code>DataLoader</code> <code>iterator</code> will recreate the <code>Dataset</code> on each worker. If you want to avoid this behavior you could use <code>persistent_workers=True</code> in the <code>DataLoader</code>, which will reuse the workers in each epoch.</p>

sxlllslgh: <p>It works! Thank you very much! <img src="https://discuss.pytorch.org/images/emoji/apple/grinning.png?v=12" title=":grinning:" class="emoji" alt=":grinning:" loading="lazy" width="20" height="20"></p>