KeyError: Caught KeyError in DataLoader worker process 0.

Hi All,

Firstly - this looks like a fab piece of work. Thank you for sharing!

I am having a stab at getting training running, and have downloaded the VGGFace2 Dataset. Folder structure looks as follows when extracted:

I point the training script at this directory and run the training script with the following command:

"python train.py --run_name first_try --dataset_path D:/projects_D/2023/ghost/train --vgg True --pretrained False --batch_size 8 --show_step 100"

The script gets the ball rolling, everythings loads and the process begins, but eventually hits a KeyError in the data loader. Full error:

Traceback (most recent call last):
File "D:\projects_D\2023\ghost\train.py", line 335, in
main(args)
File "D:\projects_D\2023\ghost\train.py", line 252, in main
train(args, device=device)
File "D:\projects_D\2023\ghost\train.py", line 232, in train
train_one_epoch(G,
File "D:\projects_D\2023\ghost\train.py", line 48, in train_one_epoch
for iteration, data in enumerate(dataloader):
File "C:\ProgramData\Anaconda3\envs\ghost_Swap\lib\site-packages\torch\utils\data\dataloader.py", line 652, in __next__
data = self._next_data()
File "C:\ProgramData\Anaconda3\envs\ghost_Swap\lib\site-packages\torch\utils\data\dataloader.py", line 1347, in _next_data
Image path: D:/projects_D/2023/ghost/train\n002097\0108_04.jpg
D:/projects_D/2023/ghost
return self._process_data(data)
File "C:\ProgramData\Anaconda3\envs\ghost_Swap\lib\site-packages\torch\utils\data\dataloader.py", line 1373, in _process_data
data.reraise()
File "C:\ProgramData\Anaconda3\envs\ghost_Swap\lib\site-packages\torch\_utils.py", line 461, in reraise
Image path: D:/projects_D/2023/ghost/train\n002447\0247_01.jpg
raise exception
**KeyError: Caught KeyError in DataLoader worker process 0.**
Original Traceback (most recent call last):
File "C:\ProgramData\Anaconda3\envs\ghost_Swap\lib\site-packages\torch\utils\data\_utils\worker.py", line 302, in _worker_loop
data = fetcher.fetch(index)
File "C:\ProgramData\Anaconda3\envs\ghost_Swap\lib\site-packages\torch\utils\data\_utils\fetch.py", line 49, in fetch
data = [self.dataset[idx] for idx in possibly_batched_index]
File "C:\ProgramData\Anaconda3\envs\ghost_Swap\lib\site-packages\torch\utils\data\_utils\fetch.py", line 49, in
data = [self.dataset[idx] for idx in possibly_batched_index]
File "D:\projects_D\2023\ghost\utils\training\Dataset.py", line 131, in __getitem__
image_path = random.choice(self.folder2imgs[folder_name])

Any bright ideas or assistance in how to resolve this would be much appreciated! As the issue lies with the self.folder2imgs dictionary, I have been trying to find a workaround without success.