balanced_sampler

data.batch_sampler.balanced_sampler

BatchSampler

BatchSampler(dataset: Dataset, batch_size: int = 1, repeat: bool = False, shuffle: bool = False, drop_last: bool = False)

Bases: Registerable, ABC

Baseclass for all batch_sampler classes.

Attributes:

• dataset (Dataset) –

  dataset to sample from.

• batch_size (int) –

  batch size.

• repeat (bool) –

  repeat data when data is not enough for batch size.

• shuffle (bool) –

  shuffle data.

• drop_last (bool) –

  drop last data when data is not enough for batch size.

Note

• batch_size, shuffle, and drop_last should be popped from dataloader config, then passed to the batch sampler config.
• repeat: for SaigeVision2 api, should be set as True.
  • if number of data is enough for batch size, then repeat is neglected.
• drop_last: for SaigeVision2 api, should be set as True.
  • to prevent the last batch from being smaller than the specified batch size.
  • if False, even when repeat is set as True, the last batch will be smaller than the specified batch size.
• for Training, shuffle should be set as True.
  • drop_last and repeat are optional.
• for Validation, repeat, shuffle, and drop_last should be set as False.

Examples:

cfg_sampler.update(
    {
        "shuffle": cfg_dataloader.pop("shuffle", False),
        "drop_last": cfg_dataloader.pop("drop_last", False),
        "batch_size": cfg_dataloader.pop("batch_size", 1),
    }
)
batch_sampler = build_batch_sampler(dataset=dataset, **cfg_sampler)
dataloader = build_dataloader(
    dataset,
    collate_fn=collate_fn,
    batch_sampler=batch_sampler,
    **cfg_dataloader,
)

Source code in SaigeToolkit/data/batch_sampler/base_sampler.py

def __init__(
    self,
    dataset: data.Dataset,
    batch_size: int = 1,
    repeat: bool = False,
    shuffle: bool = False,
    drop_last: bool = False,
) -> None:
    self.dataset = dataset

    self.batch_size = batch_size
    self.repeat = repeat

    self.shuffle = shuffle
    self.drop_last = drop_last

    self._n_repeat = None
    self._buffer_length = None

registry instance-attribute

registry: Dict[str, BatchSampler]

dataset instance-attribute

dataset = dataset

batch_size instance-attribute

batch_size = batch_size

repeat instance-attribute

repeat = repeat

shuffle instance-attribute

shuffle = shuffle

drop_last instance-attribute

drop_last = drop_last

_n_repeat instance-attribute

_n_repeat = None

_buffer_length instance-attribute

_buffer_length = None

buffer_length property

buffer_length: int

n_repeat property

n_repeat: int

__iter__

__iter__() -> Iterator[List[int]]

Source code in SaigeToolkit/data/batch_sampler/base_sampler.py

def __iter__(self) -> Iterator[List[int]]:
    if self.batch_size > self.buffer_length:
        if self.repeat:
            logger.info("data repeat activated.")
        elif self.drop_last:
            self.drop_last = False
            logger.warning("data too short. drop_last deactivated.")

    while True:
        yield from self._get_processed_buffer()

_get_processed_buffer

_get_processed_buffer() -> deque

Source code in SaigeToolkit/data/batch_sampler/base_sampler.py

def _get_processed_buffer(self) -> deque:
    buffer = self._shuffle_and_repeat_buffer()
    buffer = self._batch_buffer(buffer)
    return buffer

_shuffle_and_repeat_buffer

_shuffle_and_repeat_buffer() -> deque

Source code in SaigeToolkit/data/batch_sampler/base_sampler.py

def _shuffle_and_repeat_buffer(self) -> deque:
    if self.repeat and self.batch_size > self.buffer_length:
        buffer = deque()
        for _ in range(self.n_repeat):
            _buffer = self.generate_buffer()
            if self.shuffle:
                np.random.shuffle(_buffer)
            buffer.extend(_buffer)

    else:
        buffer = self.generate_buffer()
        if self.shuffle:
            np.random.shuffle(buffer)

    return buffer

_batch_buffer

_batch_buffer(buffer: deque) -> deque

Source code in SaigeToolkit/data/batch_sampler/base_sampler.py

def _batch_buffer(self, buffer: deque) -> deque:
    data_length, remainders = divmod(len(buffer), self.batch_size)
    _buffer = deque([[buffer.popleft() for _ in range(self.batch_size)] for _ in range(data_length)])
    if not self.drop_last and remainders > 0:
        _buffer.append([buffer.popleft() for _ in range(remainders)])
    return _buffer

__len__

__len__() -> int

Source code in SaigeToolkit/data/batch_sampler/base_sampler.py

def __len__(self) -> int:
    if self.drop_last:
        length = self.buffer_length // self.batch_size
    else:
        length = ceil(self.buffer_length / self.batch_size)
    return max(1, length)

generate_buffer abstractmethod

generate_buffer() -> deque

Source code in SaigeToolkit/data/batch_sampler/base_sampler.py

@abc.abstractmethod
def generate_buffer(self) -> deque:
    return deque([])

BalancedSampler

BalancedSampler(dataset: Dataset, balance_keys: Optional[List[str]] = None, balance_ratios: Optional[List[float]] = None, calculate_patch_number: bool = False, shuffle: bool = True, **kwargs)

Bases: BatchSampler

Sampler implementation for weighted API dataset balancing.

Sampler class is used to specify the sequence of indices/keys used in data loading. This class represents iterable objects over the indices to datasets.

When __iter__ is called, Sampler creates self.buffer, a sequence of data indices which is balanced to has exact same number of data sample from each data subset.

If balance_ratio is given, sampler autmatically re-weights each subset. Otherwise, each subset will be balanced to be equal size.

Attributes:

• dataset (Dataset) –

  dataset instance to be sampled

• calculate_patch_number (bool) –

  whether update number of patches in dataset, after sampling

• subset_data_indices (List[List[int]]) –

  split data indices, based on subset descriminator

• multipliers (ndarray) –

  number of repeats for each data subset

• remainers (ndarray) –

  number of remainers for each data subset

• shuffle (bool) –

  whether shuffle output data indices after sampling

• buffer (ndarray) –

  data indices for one epoch (accept duplication)

Note

Given:

• dataset: Dataset with number of sample 300. which can be splitted into two subsets using given key set [key_a, key_b].

  • subset_a: Data subset with number of sample 100, corresponds to key_a.
  • subset_b: Data subset with number of sample 200, corresponds to key_b.

  The word corresponds means that each domain, source, and category value parsed from the key is same to that of metadata information of each data in subset.

Output:

• iterator: An iterator with 200 indices from subset_a and 200 indices from subset_b. The max element of iterator will be 299, which is len(subset_a) + len(subset_b) - 1.

The iterator must be directly fed to torch.utils.data.DataLoader class initializer with name sample, along with torch.utils.data.ApiDataset object.

Examples:

loader = torch.utils.data.DataLoader(
   dataset=torch.utils.data.ApiDataset(...),
   shuffle=False,
   batch_sampler=BalancedDatasetSampler(
            dataset,
            balance_keys=["Document/AI_HUB", "Document/CJ"],
            balance_ratios=[0.6, 0.3],
        ),
)

Reference: https://pytorch.org/docs/stable/data.html#data-loading-order-and-sampler

initializing balance sampler

Parameters:

• dataset (Dataset) –

  dataset to be sampled. directly set to class attribute.

• balance_keys (List[str], default: None ) –

  domain-source-categoty keys as subset descriminator. Defaults to None.

• balance_ratios (Optional[List[float]], default: None ) –

  balance of subsets. if None, subsets are equally sampled. Defaults to None.

• calculate_patch_number (bool, default: False ) –

  boolean to calculate patch_number after sampling. directly set to class attribute. Defaults to False.

• shuffle (bool, default: True ) –

  boolean to shuffle data sample. Defaults to True.

Source code in SaigeToolkit/data/batch_sampler/balanced_sampler.py

def __init__(
    self,
    dataset: Dataset,
    balance_keys: Optional[List[str]] = None,
    balance_ratios: Optional[List[float]] = None,
    calculate_patch_number: bool = False,
    shuffle: bool = True,
    **kwargs,
) -> None:
    """initializing balance sampler

    Args:
        dataset (data.Dataset): dataset to be sampled. directly set to class attribute.
        balance_keys (List[str], optional):
            domain-source-categoty keys as subset descriminator. Defaults to None.
        balance_ratios (Optional[List[float]], optional): balance of subsets.
            if None, subsets are equally sampled. Defaults to None.
        calculate_patch_number (bool, optional): boolean to calculate patch_number after sampling.
            directly set to class attribute. Defaults to False.
        shuffle (bool, optional): boolean to shuffle data sample. Defaults to True.
    """

    assert hasattr(dataset, "files"), "Given dataset has no attribute 'files'!"
    assert isinstance(dataset.files, list)

    super().__init__(
        dataset,
        shuffle=shuffle,
        **kwargs,
    )
    self.calculate_patch_number = True if dataset.crop_fn is not None else calculate_patch_number
    if self.calculate_patch_number:
        assert hasattr(
            dataset, "calculate_patch_number"
        ), "function 'calculate_patch_number' should exist in dataset class"

    if balance_keys is None:
        print("No balance key provided. All domain-source-category pair will be grouped separately")
        balance_keys = self._build_default_balance_keys(dataset)
        if balance_ratios is not None:
            balance_ratios = balance_ratios if len(balance_keys) == len(balance_ratios) else None

    if balance_ratios is not None:
        assert len(balance_keys) == len(balance_ratios)
        balance_ratios = (np.array(balance_ratios) / min(balance_ratios)).tolist()
    else:
        balance_ratios = [1 for _ in range(len(balance_keys))]

    self.balance_state = {k: {"ratio": v} for k, v in zip(balance_keys, balance_ratios)}

    self._split_dataset_indices_using_keys()
    self._calculate_subset_weights()

calculate_patch_number instance-attribute

calculate_patch_number = True if crop_fn is not None else calculate_patch_number

balance_state instance-attribute

balance_state = {k: {'ratio': v} for k, v in (zip(balance_keys, balance_ratios))}

_build_default_balance_keys

_build_default_balance_keys(dataset: Dataset) -> List[str]

Build default balance_keys which is grouping all domain-source-category separately

Parameters:

• dataset (Dataset) –

  dataset instance. should have files attribute (SaigeDataset). each file data should have domain, source, category information, which would be used as subset descriminator.

Returns:

• List[str] –

  List[str]: balance_key list example: ["Document/AI_HUB", "Document/CJ"]

Author

Sukho Yoon

Source code in SaigeToolkit/data/batch_sampler/balanced_sampler.py

def _build_default_balance_keys(self, dataset: Dataset) -> List[str]:
    """Build default balance_keys which is grouping all domain-source-category separately

    Args:
        dataset (data.Dataset): dataset instance. should have `files` attribute (SaigeDataset).
            each file data should have domain, source, category information,
            which would be used as subset descriminator.

    Returns:
        List[str]: balance_key list
            example: ["Document/AI_HUB", "Document/CJ"]

    Author:
        Sukho Yoon
    """
    balance_keys = []
    for file in dataset.files:
        identity = "/".join([file["domain"], file["source"], file["category"]])
        if identity not in balance_keys:
            balance_keys.append(identity)

    return balance_keys

_extract_domain_source_category_from_key

_extract_domain_source_category_from_key(key: str) -> List[List[str]]

A function that parses given key for balancing.

empty string between slash (/) or "None" string will be treated as None-key

This function assumes target key is in the form of ...

"{domain}" # or ...
"{domain}/{source}" # or ...
"{domain}/{source}/{category}"
"''/None/{category}"

(ex) "Document/AI_HUB/상.하수도관리"

Parameters:

• key (str) –

  Key string for separating datasets into subsets.

Returns:

• List[List[str]] –

  List[List[str]]: domain_source_category_list. Key string parsed into list.

• List[List[str]] –

  (ex)

• List[List[str]] –

  ```

• List[List[str]] –

  ["Document", None, None]

• List[List[str]] –

  ["Document", "AI_HUB", None]

• List[List[str]] –

  ["Document", "AI_HUB", "상.하수도관리"]

• List[List[str]] –

  [None, None, "상.하수도관리"]

• List[List[str]] –

  ```

Author

Jonghyuk Baek

Source code in SaigeToolkit/data/batch_sampler/balanced_sampler.py

def _extract_domain_source_category_from_key(self, key: str) -> List[List[str]]:
    """
    A function that parses given key for balancing.

    empty string between slash (/) or "None" string will be treated as None-key

    This function assumes target key is in the form of ...
    ```
    "{domain}" # or ...
    "{domain}/{source}" # or ...
    "{domain}/{source}/{category}"
    "''/None/{category}"
    ```
    (ex) "Document/AI_HUB/상.하수도관리"

    Args:
        key (str): Key string for separating datasets into subsets.

        (ex)
        ```
        "Document"
        "Document/AI_HUB"
        "Document/AI_HUB/상.하수도관리"
        "/None/상.하수도관리"
        ```

    Returns:
        List[List[str]]: domain_source_category_list. Key string parsed into list.

        (ex)
        ```
        ["Document", None, None]
        ["Document", "AI_HUB", None]
        ["Document", "AI_HUB", "상.하수도관리"]
        [None, None, "상.하수도관리"]
        ```

    Author:
        Jonghyuk Baek
    """

    domain_source_category = key.split("/")
    assert (
        len(domain_source_category) <= 3
    ), f"Proper key format is {{domain}}/{{souce}}/{{category}}."
    domain_source_category += [None] * (3 - len(domain_source_category))

    return [None if key in {"", "None"} else key for key in domain_source_category]

_split_dataset_indices_using_keys

_split_dataset_indices_using_keys()

A function that splits dataset using given keys.

Each key contains metadata info: domain, source, and category. This key information will be used to separate total dataset into specific subsets.

For each data in dataset, check whether data's domain, source, and category corresponds to the key information.

For example, if a key is given as below:

key = "sample_domain/sample_source/sample_category"

Then the subset will be generated containing only index of data with metadata of:

{
    domain: "sample_domain",
    source: "sample_source",
    category: "sample_category",
    ...
}

For another example, if a key is given as below:

key = "sample_domain"

Then the subset will be generated containing only index of data with metadata of:

{
    domain: "sample_domain",
    ...
}

This generation process will be repeated for each key in target_keys from self.balance_state: target_keys: List key for subset balancing.

data_indices will be updated in self.balance_state[key] for each key. data_indices: subset_data_indices. List of data index list. Each element represents a subset of dataset.

Author

Jonghyuk Baek

Source code in SaigeToolkit/data/batch_sampler/balanced_sampler.py

def _split_dataset_indices_using_keys(self):
    """
    A function that splits dataset using given keys.

    Each key contains metadata info: `domain`, `source`, and `category`.
    This key information will be used to separate total dataset into specific subsets.

    For each data in dataset, check whether data's domain, source, and category corresponds
    to the key information.

    For example, if a key is given as below:
    ``` python
    key = "sample_domain/sample_source/sample_category"
    ```

    Then the subset will be generated containing only `index` of data with metadata of:
    ``` python
    {
        domain: "sample_domain",
        source: "sample_source",
        category: "sample_category",
        ...
    }
    ```

    For another example, if a key is given as below:
    ``` python
    key = "sample_domain"
    ```

    Then the subset will be generated containing only `index` of data with metadata of:
    ``` python
    {
        domain: "sample_domain",
        ...
    }
    ```

    This generation process will be repeated for each key in `target_keys` from self.balance_state:
        target_keys: List key for subset balancing.

    `data_indices` will be updated in self.balance_state[key] for each key.
        data_indices: subset_data_indices. List of data index list.
            Each element represents a subset of dataset.

    Author:
        Jonghyuk Baek
    """

    # To filter out concateneted dataset into subsets.
    target_keys = self.balance_state.keys()

    target_domain_source_category_list = []

    for target_key in target_keys:
        domain_source_category = self._extract_domain_source_category_from_key(target_key)
        target_domain_source_category_list.append(domain_source_category)

    # Split dataset into subsets using balancing target keys.
    # self.dataset.files: List containing each data's metadata dictionary.
    subset_data_indices = [[] for _ in range(len(target_domain_source_category_list))]
    for idx_data, metadata_dict in enumerate(self.dataset.files):
        data_domain_source_category = [metadata_dict[k] for k in ["domain", "source", "category"]]

        # Fill out each subset with matching condition.
        for idx_subset, target_domain_source_category in enumerate(
            target_domain_source_category_list
        ):
            if self.match_identifier(data_domain_source_category, target_domain_source_category):
                subset_data_indices[idx_subset].append(idx_data)
                break

    for key, indices in zip(target_keys, subset_data_indices):
        assert (
            len(indices) != 0
        ), f"No matching data for key [{key}]. Please check balance target keys.\n"
        self.balance_state[key].update({"data_indices": indices})

match_identifier

match_identifier(id_data: List[Optional[str]], id_target: List[Optional[str]]) -> bool

Check data identifier matches non-None target identifier identifier: [domain, source, category]

Parameters:

• id_data (Optional[str]) –

  data item identifier

• id_target (Optional[str]) –

  balance target subset identifier

Returns:

• bool ( bool ) –

  whether id's are matched

Author

Sukho Yoon

Source code in SaigeToolkit/data/batch_sampler/balanced_sampler.py

def match_identifier(self, id_data: List[Optional[str]], id_target: List[Optional[str]]) -> bool:
    """Check data identifier matches non-None target identifier
    identifier: [domain, source, category]

    Args:
        id_data (Optional[str]): data item identifier
        id_target (Optional[str]): balance target subset identifier

    Returns:
        bool: whether id's are matched

    Author:
        Sukho Yoon
    """

    for id_d, id_t in zip(id_data, id_target):
        if id_t is None:
            continue

        if id_d != id_t:
            return False

    return True

_update_subset_sizes

_update_subset_sizes()

Source code in SaigeToolkit/data/batch_sampler/balanced_sampler.py

def _update_subset_sizes(self):
    for key, balance_state in self.balance_state.items():
        indices = balance_state["data_indices"]
        if self.calculate_patch_number:
            num_patch = self._update_dataset_num_patch(indices)
            self.balance_state[key].update(
                {"size": num_patch, "patch_per_image": num_patch / len(indices)}
            )
        else:
            self.balance_state[key].update({"size": len(indices)})

_calculate_subset_weights

_calculate_subset_weights()

Calculate subset multiplier and remainer using following two variables from self.balance_state: balance_ratios: resultant data exposure ratio by subsets subset_size_ratios: data sizes ratio between subsets

Author

Sukho Yoon

Source code in SaigeToolkit/data/batch_sampler/balanced_sampler.py

def _calculate_subset_weights(self):
    """
    Calculate subset multiplier and remainer using following two variables from self.balance_state:
        balance_ratios: resultant data exposure ratio by subsets
        subset_size_ratios: data sizes ratio between subsets

    Author:
        Sukho Yoon
    """
    self._update_subset_sizes()

    balance_ratios = np.array([v["ratio"] for v in self.balance_state.values()])
    subset_sizes = np.array([v["size"] for v in self.balance_state.values()])
    subset_repeatance = balance_ratios / subset_sizes
    subset_repeatance = subset_repeatance / min(subset_repeatance)

    target_sizes = np.round(subset_repeatance * subset_sizes).astype(int)
    multipliers = np.floor(subset_repeatance).astype(int)
    remainers = target_sizes - multipliers * subset_sizes

    for key, multiplier, remainer in zip(self.balance_state.keys(), multipliers, remainers):
        if self.calculate_patch_number:
            remainer = np.round(remainer / self.balance_state[key]["patch_per_image"]).astype(int)

        self.balance_state[key].update({"multiplier": multiplier, "remainer": remainer})

generate_buffer

generate_buffer() -> deque

Creates index list of data instances, which is already balanced.

Exact result can be always different even when shuffle=False when non-zero element in self.balance_state[key]["remainer"] exists.

Author

Jonghyuk Baek

Source code in SaigeToolkit/data/batch_sampler/balanced_sampler.py

def generate_buffer(self) -> deque:
    """
    Creates index list of data instances, which is already balanced.

    Exact result can be always different even when `shuffle=False` when non-zero element
    in `self.balance_state[key]["remainer"]` exists.

    Author:
        Jonghyuk Baek
    """

    buffer = []

    if self.calculate_patch_number:
        n_patch = 0

    sample_with_replacement = bool(self.calculate_patch_number)
    # Expand data index list from smaller dataset, to be same length as the biggest one.
    for balance_state in self.balance_state.values():
        indices = balance_state["data_indices"]
        multiplier = balance_state["multiplier"]
        remainer = balance_state["remainer"]

        # update multipliers
        buffer.append(np.repeat(indices, multiplier))
        if self.calculate_patch_number:
            n_patch += balance_state["size"] * multiplier

        # update remainers
        _buffer_remainers = np.random.choice(indices, remainer, replace=sample_with_replacement)
        buffer.append(_buffer_remainers)
        if self.calculate_patch_number:
            n_patch += self._update_dataset_num_patch(_buffer_remainers)

    if self.calculate_patch_number:
        self.dataset.n_patch = n_patch

    # Concatenate index list from every dataset and shuffle.
    # Shuffle parameter should not be fed when initializing `torch.utils.data.DataLoader`.
    buffer = np.concatenate(buffer, axis=None).tolist()
    buffer = deque(buffer)

    return buffer

_update_dataset_num_patch

_update_dataset_num_patch(indices: Union[deque, List[int]]) -> int

Updates n_patch attribute in dataset, when indices is updated. (due to changes in remainer items)

Parameters:

• indices (Union[deque, List[int]]) –

  file indices for calculating data patch number

Returns:

• int ( int ) –

  number of patches in files

Author

Sukho Yoon

Source code in SaigeToolkit/data/batch_sampler/balanced_sampler.py

def _update_dataset_num_patch(self, indices: Union[deque, List[int]]) -> int:
    """Updates n_patch attribute in dataset, when indices is updated.
    (due to changes in remainer items)

    Args:
        indices (Union[deque, List[int]]): file indices for calculating data patch number

    Returns:
        int: number of patches in files

    Author:
        Sukho Yoon
    """

    subset_files = [self.dataset.files[index] for index in indices]
    return self.dataset.calculate_patch_number(subset_files)