Блог

Мой набор данных зависит от тензора 3GB. Этот тензор может быть либо на CPU, либо на GPU. Узким местом моего кода является предварительная обработка загрузки данных. Но я не могу добавить больше нескольких рабочих, не убивая свою оперативную память.

Для меня это звучит глупо: почему каждый работник может получить копию тензора объемом 3 ГБ, когда он точно такой же для каждого работника?

Есть ли какое-либо решение для предоставления рабочим доступа к одной версии этого тензора?

Спасибо,

В документации Pytorch явно упоминается эта проблема с DataLoader дублированием базового набора данных (по крайней мере, в Windows и macOS, насколько я понимаю).

В общем, вам не следует с нетерпением загружать весь ваш набор данных в память из-за такой проблемы. Набор данных должен загружаться с задержкой, т. Е. Образцы должны загружаться только при обращении к ним в __getitem__ методе.

Если весь ваш набор данных хранится на диске в виде монолитного тензора, вы можете разбить его на отдельные выборки и сохранить их, например, в папке.

Затем вы могли бы определить свой набор данных как:

 from torch.utils.data import Dataset, DataLoader
from glob import glob
from os.path import abspath


class MyDataset(Dataset):
  def __init__(self, folder: str):
    # Retrieve all tensor file names
    folder = abspath(folder)
    self.files = glob(f"{folder}/*.pt")

  def __getitem__(self, index: int):
    # Load tensors on demand
    return torch.load(self.files[index])
 
  def __len__(self) -> int:
    return len(self.files)
  

Другим решением является сопоставление данных с памятью. Это то, что HuggingFace делает для огромных наборов данных, посмотрите здесь. Это позволяет избежать загрузки всего набора данных в ОЗУ, а также позволяет совместно использовать его в нескольких процессах без копирования.

Ray может быть интересным вариантом для вас. Проверьте наборы данных для обучения лучей!

Кроме того, вы также можете использовать

 data_id = ray.put(data)
  

чтобы сбросить ваши данные, и

 data = ray.get(data_id)
  

чтобы загружать одни и те же файлы, не копируя их между функциями.