fix: fix error with multiprocessing in multigpu filter

Author: boomb0om

DPF/filters/data_filter.py

@@ -1,3 +1,5 @@
+import multiprocessing
+import sys
 from abc import ABC, abstractmethod
 from typing import Any
 
@@ -19,6 +21,7 @@ def __init__(self, pbar: bool, _pbar_position: int = 0):
  super().__init__()
  self.pbar = pbar
  self.pbar_position = _pbar_position
+ self._created_by_multigpu_data_filter = False
 
  @property
  def schema(self) -> list[str]:
@@ -120,7 +123,15 @@ def run(self, dataset: Dataset[tuple[bool, Any]]) -> pd.DataFrame:
  pd.DataFrame
  Dataframe with columns from schema property
  """
- dataloader = DataLoader(dataset, collate_fn=identical_collate_fn, **self.dataloader_kwargs)
+ multiprocessing_context = None
+ if self._created_by_multigpu_data_filter and sys.platform not in {'win32', 'darwin'}:
+ multiprocessing_context = multiprocessing.get_context('fork')
+
+ dataloader = DataLoader(
+ dataset, collate_fn=identical_collate_fn,
+ multiprocessing_context=multiprocessing_context,
+ **self.dataloader_kwargs
+ )
  filter_results = self._get_dict_from_schema()
 
  for batch in tqdm(dataloader, disable=not self.pbar, position=self.pbar_position):

DPF/filters/multigpu_filter.py

@@ -29,6 +29,7 @@ def run_one_process(
  reader = DatasetReader(connector=connector)
  processor = reader.from_df(config, df)
  datafilter = filter_class(**filter_kwargs, _pbar_position=i, device=device) # type: ignore
+ datafilter._created_by_multigpu_data_filter = True
  processor.apply_data_filter(datafilter, **filter_run_kwargs)
  res = processor.df
  res.set_index(index, inplace=True)
@@ -118,8 +119,8 @@ def run(
  )
 
  processes = []
+ context = multiprocessing.get_context('spawn')
  for param in params:
- context = multiprocessing.get_context('spawn')
  p = context.Process(target=run_one_process, args=param)
  p.start()
  processes.append(p)

DPF/filters/videos/lita_filter.py

@@ -82,6 +82,8 @@ def __init__(
  pretrainers = load_pretrained_model(weights_path, model_base, self.model_name, load_8bit, load_4bit)
  self.tokenizer, self.model, self.processor, self.context_len = pretrainers
 
+ self.model_num_frames = self.model.config.num_frames
+
  self.conv_mode = "llava_v1"
  self.conv = conv_templates[self.conv_mode].copy()
 
@@ -118,7 +120,7 @@ def preprocess_data(
  ) -> Any:
  key = metadata[self.key_column]
  video_file = BytesIO(modality2data['video'])
- video_file = load_video(video_file, self.processor, self.model.config.num_frames).unsqueeze(0).half()
+ video_file = load_video(video_file, self.processor, self.model_num_frames).unsqueeze(0).half()
  return key, video_file
 
  def process_batch(self, batch: list[Any]) -> dict[str, list[Any]]: