Shortcuts

Source code for ding.data.buffer.middleware.use_time_check

from typing import Callable, Any, List, Optional, Union, TYPE_CHECKING
from collections import defaultdict
from ding.data.buffer import BufferedData
if TYPE_CHECKING:
    from ding.data.buffer.buffer import Buffer


[docs]def use_time_check(buffer_: 'Buffer', max_use: int = float("inf")) -> Callable: """ Overview: This middleware aims to check the usage times of data in buffer. If the usage times of a data is greater than or equal to max_use, this data will be removed from buffer as soon as possible. Arguments: - max_use (:obj:`int`): The max reused (resampled) count for any individual object. """ use_count = defaultdict(int) def _need_delete(item: BufferedData) -> bool: nonlocal use_count idx = item.index use_count[idx] += 1 item.meta['use_count'] = use_count[idx] if use_count[idx] >= max_use: return True else: return False def _check_use_count(sampled_data: List[BufferedData]): delete_indices = [item.index for item in filter(_need_delete, sampled_data)] buffer_.delete(delete_indices) for index in delete_indices: del use_count[index] def sample(chain: Callable, *args, **kwargs) -> Union[List[BufferedData], List[List[BufferedData]]]: sampled_data = chain(*args, **kwargs) if len(sampled_data) == 0: return sampled_data if isinstance(sampled_data[0], BufferedData): _check_use_count(sampled_data) else: for grouped_data in sampled_data: _check_use_count(grouped_data) return sampled_data def _use_time_check(action: str, chain: Callable, *args, **kwargs) -> Any: if action == "sample": return sample(chain, *args, **kwargs) return chain(*args, **kwargs) return _use_time_check