Skip to content

dataloader

Module diagram

classDiagram
  class dataloader {
  }
  class builder {
  }
  class full2patch_loader {
  }
  class infinite_random_sampler {
  }
  dataloader --> builder
  builder --> full2patch_loader

data.dataloader

SaigeDataLoader module-attribute 

SaigeDataLoader = Union[DataLoader, Full2PatchDataLoader]

get_dataloader_class 

get_dataloader_class(cfg_dataloader_name: Optional[str] = None, crop_fn: Optional[Callable[[dict], List[dict]]] = None) -> Type[SaigeDataLoader]

getting dataloader class

Parameters:

  • cfg_dataloader_name (Optional[str], default: None ) –

    name of dataloader class. If None, choose appropriate class considering crop_fn. Defaults to None.

  • crop_fn (Optional[Callable[[dict], List[dict]]], default: None ) –

    image crop function. Used for segmentation/OCR patch-wise training. Defaults to None.

Returns:

Source code in SaigeToolkit/data/dataloader/builder.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def get_dataloader_class(
    cfg_dataloader_name: Optional[str] = None, crop_fn: Optional[Callable[[dict], List[dict]]] = None
) -> Type[SaigeDataLoader]:
    """getting dataloader class

    Args:
        cfg_dataloader_name (Optional[str], optional): name of dataloader class.
            If None, choose appropriate class considering `crop_fn`. Defaults to None.
        crop_fn (Optional[Callable[[dict], List[dict]]], optional): image crop function.
            Used for segmentation/OCR patch-wise training. Defaults to None.

    Returns:
        Type[SaigeDataLoader]: dataloader class
    """
    if cfg_dataloader_name:
        try:
            return {
                "torch": data.DataLoader,
                "full2patch": Full2PatchDataLoader,
            }[cfg_dataloader_name]
        except:
            raise (f"Dataset {cfg_dataloader_name} not available")

    if crop_fn:
        return Full2PatchDataLoader

    # monkey patch
    def _stop(self):
        pass

    data.DataLoader.stop = _stop

    return data.DataLoader

builder

get_dataloader_class 

get_dataloader_class(cfg_dataloader_name: Optional[str] = None, crop_fn: Optional[Callable[[dict], List[dict]]] = None) -> Type[SaigeDataLoader]

getting dataloader class

Parameters:

  • cfg_dataloader_name (Optional[str], default: None ) –

    name of dataloader class. If None, choose appropriate class considering crop_fn. Defaults to None.

  • crop_fn (Optional[Callable[[dict], List[dict]]], default: None ) –

    image crop function. Used for segmentation/OCR patch-wise training. Defaults to None.

Returns:

Source code in SaigeToolkit/data/dataloader/builder.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def get_dataloader_class(
    cfg_dataloader_name: Optional[str] = None, crop_fn: Optional[Callable[[dict], List[dict]]] = None
) -> Type[SaigeDataLoader]:
    """getting dataloader class

    Args:
        cfg_dataloader_name (Optional[str], optional): name of dataloader class.
            If None, choose appropriate class considering `crop_fn`. Defaults to None.
        crop_fn (Optional[Callable[[dict], List[dict]]], optional): image crop function.
            Used for segmentation/OCR patch-wise training. Defaults to None.

    Returns:
        Type[SaigeDataLoader]: dataloader class
    """
    if cfg_dataloader_name:
        try:
            return {
                "torch": data.DataLoader,
                "full2patch": Full2PatchDataLoader,
            }[cfg_dataloader_name]
        except:
            raise (f"Dataset {cfg_dataloader_name} not available")

    if crop_fn:
        return Full2PatchDataLoader

    # monkey patch
    def _stop(self):
        pass

    data.DataLoader.stop = _stop

    return data.DataLoader

full2patch_loader

Full2PatchDataLoader 

Full2PatchDataLoader(dataset: Dataset, collate_fn: Optional[Callable[[List[dict]], dict]] = None, batch_sampler: Optional[Iterable[int]] = None, batch_size: int = 1, shuffle: bool = False, drop_last: bool = True, **cfg_dataloader: dict)

DataLoader class for Segmentation and SceneTextRecognition tasks. Whole image is divided into multiple image patches, based on polygon labels. Full2PatchDataLoader has dataloader as class attribute, which returns croppped image patches. Image patches are stocked in buffer attribute, then passed to deep network in batch_size.

Attributes:

  • dataset (Dataset) –

    base dataset

  • batch_size (int) –

    batch size

  • drop_last (bool) –

    whether drop remaining data items less than batch_size at last iter.

  • buffer_multiple (int) –

    maximum size of buffer is decided by {buffer_multiple} x {batch_size}.

  • loader (DataLoader) –

    DataLoader that provides cropped image patches.

  • loader_iter (Iterable[list]) –

    iterator of loader attribute.

  • buffer (List[dict]) –

    buffer where image patches are stacking.

  • collate_function (Callable[[List[dict]], dict]) –

    function collating image patches into a batch.

initializing Full2PatchLoader

Parameters:

  • dataset (Dataset) –

    base dataset to be loaded.

  • collate_fn (Optional[Callable[[List[dict]], dict]], default: None ) –

    function collating image patches into a batch. Defaults to None.

  • sampler (Optional[Iterable[int]]) –

    specific sampler such as balanced sampler. Defaults to None.

  • batch_size (int, default: 1 ) –

    size of a batch for one iteration. Defaults to 1.

  • drop_last (bool, default: True ) –

    drop_last boolean. Defaults to True.

Source code in SaigeToolkit/data/dataloader/full2patch_loader.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
def __init__(
    self,
    dataset: data.Dataset,
    collate_fn: Optional[Callable[[List[dict]], dict]] = None,
    batch_sampler: Optional[Iterable[int]] = None,
    batch_size: int = 1,
    shuffle: bool = False,
    drop_last: bool = True,
    **cfg_dataloader: dict,
):
    """initializing Full2PatchLoader

    Args:
        dataset (data.Dataset): base dataset to be loaded.
        collate_fn (Optional[Callable[[List[dict]], dict]], optional):
            function collating image patches into a batch. Defaults to None.
        sampler (Optional[Iterable[int]], optional):
            specific sampler such as balanced sampler. Defaults to None.
        batch_size (int, optional): size of a batch for one iteration. Defaults to 1.
        drop_last (bool, optional): drop_last boolean. Defaults to True.
    """

    self.dataset = dataset
    self.batch_size = batch_size
    self.shuffle = shuffle
    self.drop_last = drop_last

    self.buffer_multiple = 10
    self.loader = data.DataLoader(
        dataset,
        collate_fn=self.collate_patch,
        batch_sampler=batch_sampler,
        **cfg_dataloader,
    )
    self.loader_iter = None
    self.buffer = []

    self.collate_function = collate_fn if collate_fn is not None else self.collate_batch
collate_patch 
collate_patch(batch: List[List[dict]]) -> List[dict]

dataset with crop_fn returns List[List[dict]] type. Since batch_size of self.loader is hard-defined as 1, simply returning first item of input is enough.

Parameters:

  • batch (List[List[dict]]) –

    batch (1) of cropped patches (n)

Returns:

  • List[dict] –

    List[dict]: cropped patches (n)

Source code in SaigeToolkit/data/dataloader/full2patch_loader.py 
68
69
70
71
72
73
74
75
76
77
78
79
def collate_patch(self, batch: List[List[dict]]) -> List[dict]:
    """dataset with crop_fn returns List[List[dict]] type.
    Since batch_size of self.loader is hard-defined as 1,
    simply returning first item of input is enough.

    Args:
        batch (List[List[dict]]): batch (1) of cropped patches (n)

    Returns:
        List[dict]: cropped patches (n)
    """
    return batch[0]
collate_batch 
collate_batch(batch: List[dict]) -> dict

collating image patches into a batch

Parameters:

  • batch (List[dict]) –

    list cropped patch data dicts (n)

Returns:

  • dict ( dict ) –

    a batch dict

Source code in SaigeToolkit/data/dataloader/full2patch_loader.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def collate_batch(self, batch: List[dict]) -> dict:
    """collating image patches into a batch

    Args:
        batch (List[dict]): list cropped patch data dicts (n)

    Returns:
        dict: a batch dict
    """

    data_out = {k: [] for k in batch[0].keys()}
    for data in batch:
        for key, value in data.items():
            if torch.is_tensor(value):
                value = value[None]
            data_out[key].append(value)

    for key, value in data_out.items():
        if torch.is_tensor(value[0]):
            try:
                data_out[key] = torch.cat(value)
            except:
                print(f'[{"WARNING".center(9)}] {key} is not concatenated')

    return data_out
__next__ 
__next__() -> dict

generating a batch data, extracting from buffer stack.

Raises:

  • StopIteration –

    self.get_buffer() called due to (len(self.buffer) < self.batch_size) and (len(self.buffer) == 0) even after self.get_buffer() : raise StopIteration.

  • StopIteration –

    self.get_buffer() called due to (len(self.buffer) < self.batch_size) and (len(self.buffer) < self.batch_size) even after self.get_buffer() and (drop_last == True) : raise StopIteration.

Returns:

  • dict ( dict ) –

    a batch dict

Source code in SaigeToolkit/data/dataloader/full2patch_loader.py 
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def __next__(self) -> dict:
    """generating a batch data, extracting from buffer stack.

    Raises:
        StopIteration:
            self.get_buffer() called due to (len(self.buffer) < self.batch_size)
            and (len(self.buffer) == 0) even after self.get_buffer()
            : raise StopIteration.
        StopIteration:
            self.get_buffer() called due to (len(self.buffer) < self.batch_size)
            and (len(self.buffer) < self.batch_size) even after self.get_buffer()
            and (drop_last == True)
            : raise StopIteration.

    Returns:
        dict: a batch dict
    """

    batch_size = self.batch_size
    if len(self.buffer) < self.batch_size:
        self.get_buffer()

        if len(self.buffer) == 0:
            raise StopIteration

        if len(self.buffer) < self.batch_size:
            if self.drop_last:
                raise StopIteration
            batch_size = len(self.buffer)

    if self.shuffle:
        idx = np.sort(np.random.choice(range(len(self.buffer)), size=batch_size, replace=False))
    else:
        idx = np.array([0 for _ in range(batch_size)])

    batched_buffer = [self.buffer.pop(idx[-i - 1]) for i in range(len(idx))]

    return self.collate_function(batched_buffer)

infinite_random_sampler

InfiniteRandomSampler

Bases: RandomSampler

Infinitely samples elements randomly. If without replacement, then sample from a shuffled dataset. If with replacement, then user can specify :attr:num_samples to draw.

Parameters:

  • data_source (Dataset) –

    dataset to sample from

  • replacement (bool) –

    samples are drawn on-demand with replacement if True, default=False

  • num_samples (int) –

    number of samples to draw, default=len(dataset).

  • generator (Generator) –

    Generator used in sampling.

Assume that this DataLoader is used for training with shuffle=True.

We use InfiniteRandomSampler to make the DataLoader to infinitely sample the dataset. This prevents the "drop_last" and the "prefeching" problem across epochs. Be careful that you should stop the training loop by counting the number of steps manually. The DataLoader will never stop by itself.

NOTE: Note that this sampler has infinite length and thus you should be careful when calculating the current epoch. We recommend you to use step//steps_per_epoch to calculate the current epoch.