Function _enqueue_batch

examples/mlperf/dataloader.py:383–396

(bc)

def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, shuffle:bool=True, seed:int|None=None):
  def _enqueue_batch(bc):
    from extra.datasets.openimages import prepare_target
    for idx in range(bc * batch_size, (bc+1) * batch_size):
      img = dataset.loadImgs(next(dataset_iter))[0]
      ann = dataset.loadAnns(dataset.getAnnIds(img_id:=img["id"]))
      tgt = prepare_target(ann, img_id, (img["height"], img["width"]))

      if img_ids is not None:
        img_ids[idx] = img_id

      if img_sizes is not None:
        img_sizes[idx] = tgt["image_size"]

      queue_in.put((idx, img, tgt))

  def _setup_shared_mem(shm_name:str, size:tuple[int, ...], dtype:dtypes) -> tuple[shared_memory.SharedMemory, Tensor]:
    shm_name = f"{shm_name}_{os.getpid()}"