Проблема заключается в следующем:
ds = ds.shuffle(buffer_size = len(all_image_paths))
Буфер, используемый Dataset.shuffle()
, является буфером «в памяти», поэтому вы фактически пытаетесь загрузить весь набор данных в память.
У вас есть несколько вариантов (которые вы можете комбинировать), чтобы исправить это:
Вариант 1:
Уменьшить размер буфера до гораздо меньшего числа.
Вариант 2:
Переместите shuffle()
перед оператором map()
.
Это означает, что мы будем перетасовывать до того, как загрузим изображения, поэтому мы будем просто хранить имена файлов в буфере памяти для перемешивания, а не хранить огромные тензоры.