Shortcuts

Source code for ding.worker.learner.base_learner

from typing import Any, Union, Callable, List, Dict, Optional, Tuple
from ditk import logging
from collections import namedtuple
from functools import partial
from easydict import EasyDict

import copy

from ding.torch_utils import CountVar, auto_checkpoint, build_log_buffer
from ding.utils import build_logger, EasyTimer, import_module, LEARNER_REGISTRY, get_rank, get_world_size
from ding.utils.autolog import LoggedValue, LoggedModel, TickTime
from ding.utils.data import AsyncDataLoader
from .learner_hook import build_learner_hook_by_cfg, add_learner_hook, merge_hooks, LearnerHook


[docs]@LEARNER_REGISTRY.register('base') class BaseLearner(object): r""" Overview: Base class for policy learning. Interface: train, call_hook, register_hook, save_checkpoint, start, setup_dataloader, close Property: learn_info, priority_info, last_iter, train_iter, rank, world_size, policy monitor, log_buffer, logger, tb_logger, ckpt_name, exp_name, instance_name """ @classmethod def default_config(cls: type) -> EasyDict: cfg = EasyDict(copy.deepcopy(cls.config)) cfg.cfg_type = cls.__name__ + 'Dict' return cfg config = dict( train_iterations=int(1e9), dataloader=dict(num_workers=0, ), log_policy=True, # --- Hooks --- hook=dict( load_ckpt_before_run='', log_show_after_iter=100, save_ckpt_after_iter=10000, save_ckpt_after_run=True, ), ) _name = "BaseLearner" # override this variable for sub-class learner
[docs] def __init__( self, cfg: EasyDict, policy: namedtuple = None, tb_logger: Optional['SummaryWriter'] = None, # noqa dist_info: Tuple[int, int] = None, exp_name: Optional[str] = 'default_experiment', instance_name: Optional[str] = 'learner', ) -> None: """ Overview: Initialization method, build common learner components according to cfg, such as hook, wrapper and so on. Arguments: - cfg (:obj:`EasyDict`): Learner config, you can refer cls.config for details. - policy (:obj:`namedtuple`): A collection of policy function of learn mode. And policy can also be \ initialized when runtime. - tb_logger (:obj:`SummaryWriter`): Tensorboard summary writer. - dist_info (:obj:`Tuple[int, int]`): Multi-GPU distributed training information. - exp_name (:obj:`str`): Experiment name, which is used to indicate output directory. - instance_name (:obj:`str`): Instance name, which should be unique among different learners. Notes: If you want to debug in sync CUDA mode, please add the following code at the beginning of ``__init__``. .. code:: python os.environ['CUDA_LAUNCH_BLOCKING'] = "1" # for debug async CUDA """ self._cfg = cfg self._exp_name = exp_name self._instance_name = instance_name self._ckpt_name = None self._timer = EasyTimer() # These 2 attributes are only used in parallel mode. self._end_flag = False self._learner_done = False if dist_info is None: self._rank = get_rank() self._world_size = get_world_size() else: # Learner rank. Used to discriminate which GPU it uses. self._rank, self._world_size = dist_info if self._world_size > 1: self._cfg.hook.log_reduce_after_iter = True # Logger (Monitor will be initialized in policy setter) # Only rank == 0 learner needs monitor and tb_logger, others only need text_logger to display terminal output. if self._rank == 0: if tb_logger is not None: self._logger, _ = build_logger( './{}/log/{}'.format(self._exp_name, self._instance_name), self._instance_name, need_tb=False ) self._tb_logger = tb_logger else: self._logger, self._tb_logger = build_logger( './{}/log/{}'.format(self._exp_name, self._instance_name), self._instance_name ) else: self._logger, _ = build_logger( './{}/log/{}'.format(self._exp_name, self._instance_name), self._instance_name, need_tb=False ) self._tb_logger = None self._log_buffer = { 'scalar': build_log_buffer(), 'scalars': build_log_buffer(), 'histogram': build_log_buffer(), } # Setup policy if policy is not None: self.policy = policy # Learner hooks. Used to do specific things at specific time point. Will be set in ``_setup_hook`` self._hooks = {'before_run': [], 'before_iter': [], 'after_iter': [], 'after_run': []} # Last iteration. Used to record current iter. self._last_iter = CountVar(init_val=0) # Collector envstep. Used to record current envstep. self._collector_envstep = 0 # Setup time wrapper and hook. self._setup_wrapper() self._setup_hook()
[docs] def _setup_hook(self) -> None: """ Overview: Setup hook for base_learner. Hook is the way to implement some functions at specific time point in base_learner. You can refer to ``learner_hook.py``. """ if hasattr(self, '_hooks'): self._hooks = merge_hooks(self._hooks, build_learner_hook_by_cfg(self._cfg.hook)) else: self._hooks = build_learner_hook_by_cfg(self._cfg.hook)
[docs] def _setup_wrapper(self) -> None: """ Overview: Use ``_time_wrapper`` to get ``train_time``. Note: ``data_time`` is wrapped in ``setup_dataloader``. """ self._wrapper_timer = EasyTimer() self.train = self._time_wrapper(self.train, 'scalar', 'train_time')
def _time_wrapper(self, fn: Callable, var_type: str, var_name: str) -> Callable: """ Overview: Wrap a function and record the time it used in ``_log_buffer``. Arguments: - fn (:obj:`Callable`): Function to be time_wrapped. - var_type (:obj:`str`): Variable type, e.g. ['scalar', 'scalars', 'histogram']. - var_name (:obj:`str`): Variable name, e.g. ['cur_lr', 'total_loss']. Returns: - wrapper (:obj:`Callable`): The wrapper to acquire a function's time. """ def wrapper(*args, **kwargs) -> Any: with self._wrapper_timer: ret = fn(*args, **kwargs) self._log_buffer[var_type][var_name] = self._wrapper_timer.value return ret return wrapper
[docs] def register_hook(self, hook: LearnerHook) -> None: """ Overview: Add a new learner hook. Arguments: - hook (:obj:`LearnerHook`): The hook to be addedr. """ add_learner_hook(self._hooks, hook)
@property def collector_envstep(self) -> int: """ Overview: Get current collector envstep. Returns: - collector_envstep (:obj:`int`): Current collector envstep. """ return self._collector_envstep @collector_envstep.setter def collector_envstep(self, value: int) -> None: """ Overview: Set current collector envstep. Arguments: - value (:obj:`int`): Current collector envstep. """ self._collector_envstep = value
[docs] def train(self, data: dict, envstep: int = -1, policy_kwargs: Optional[dict] = None) -> None: """ Overview: Given training data, implement network update for one iteration and update related variables. Learner's API for serial entry. Also called in ``start`` for each iteration's training. Arguments: - data (:obj:`dict`): Training data which is retrieved from repaly buffer. .. note:: ``_policy`` must be set before calling this method. ``_policy.forward`` method contains: forward, backward, grad sync(if in multi-gpu mode) and parameter update. ``before_iter`` and ``after_iter`` hooks are called at the beginning and ending. """ assert hasattr(self, '_policy'), "please set learner policy" self.call_hook('before_iter') if policy_kwargs is None: policy_kwargs = {} # Forward log_vars = self._policy.forward(data, **policy_kwargs) # Update replay buffer's priority info if isinstance(log_vars, dict): priority = log_vars.pop('priority', None) elif isinstance(log_vars, list): priority = log_vars[-1].pop('priority', None) else: raise TypeError("not support type for log_vars: {}".format(type(log_vars))) if priority is not None: replay_buffer_idx = [d.get('replay_buffer_idx', None) for d in data] replay_unique_id = [d.get('replay_unique_id', None) for d in data] self.priority_info = { 'priority': priority, 'replay_buffer_idx': replay_buffer_idx, 'replay_unique_id': replay_unique_id, } # Discriminate vars in scalar, scalars and histogram type # Regard a var as scalar type by default. For scalars and histogram type, must annotate by prefix "[xxx]" self._collector_envstep = envstep if isinstance(log_vars, dict): log_vars = [log_vars] for elem in log_vars: scalars_vars, histogram_vars = {}, {} for k in list(elem.keys()): if "[scalars]" in k: new_k = k.split(']')[-1] scalars_vars[new_k] = elem.pop(k) elif "[histogram]" in k: new_k = k.split(']')[-1] histogram_vars[new_k] = elem.pop(k) # Update log_buffer self._log_buffer['scalar'].update(elem) self._log_buffer['scalars'].update(scalars_vars) self._log_buffer['histogram'].update(histogram_vars) self.call_hook('after_iter') self._last_iter.add(1) return log_vars
@auto_checkpoint def start(self) -> None: """ Overview: [Only Used In Parallel Mode] Learner's API for parallel entry. For each iteration, learner will get data through ``_next_data`` and call ``train`` to train. .. note:: ``before_run`` and ``after_run`` hooks are called at the beginning and ending. """ self._end_flag = False self._learner_done = False # before run hook self.call_hook('before_run') for i in range(self._cfg.train_iterations): data = self._next_data() if self._end_flag: break self.train(data) self._learner_done = True # after run hook self.call_hook('after_run')
[docs] def setup_dataloader(self) -> None: """ Overview: [Only Used In Parallel Mode] Setup learner's dataloader. .. note:: Only in parallel mode will we use attributes ``get_data`` and ``_dataloader`` to get data from file system; Instead, in serial version, we can fetch data from memory directly. In parallel mode, ``get_data`` is set by ``LearnerCommHelper``, and should be callable. Users don't need to know the related details if not necessary. """ cfg = self._cfg.dataloader batch_size = self._policy.get_attribute('batch_size') device = self._policy.get_attribute('device') chunk_size = cfg.chunk_size if 'chunk_size' in cfg else batch_size self._dataloader = AsyncDataLoader( self.get_data, batch_size, device, chunk_size, collate_fn=lambda x: x, num_workers=cfg.num_workers ) self._next_data = self._time_wrapper(self._next_data, 'scalar', 'data_time')
def _next_data(self) -> Any: """ Overview: [Only Used In Parallel Mode] Call ``_dataloader``'s ``__next__`` method to return next training data. Returns: - data (:obj:`Any`): Next training data from dataloader. """ return next(self._dataloader)
[docs] def close(self) -> None: """ Overview: [Only Used In Parallel Mode] Close the related resources, e.g. dataloader, tensorboard logger, etc. """ if self._end_flag: return self._end_flag = True if hasattr(self, '_dataloader'): self._dataloader.close() if self._tb_logger: self._tb_logger.flush() self._tb_logger.close()
def __del__(self) -> None: self.close()
[docs] def call_hook(self, name: str) -> None: """ Overview: Call the corresponding hook plugins according to position name. Arguments: - name (:obj:`str`): Hooks in which position to call, \ should be in ['before_run', 'after_run', 'before_iter', 'after_iter']. """ for hook in self._hooks[name]: hook(self)
def info(self, s: str) -> None: """ Overview: Log string info by ``self._logger.info``. Arguments: - s (:obj:`str`): The message to add into the logger. """ self._logger.info('[RANK{}]: {}'.format(self._rank, s)) def debug(self, s: str) -> None: self._logger.debug('[RANK{}]: {}'.format(self._rank, s))
[docs] def save_checkpoint(self, ckpt_name: str = None) -> None: """ Overview: Directly call ``save_ckpt_after_run`` hook to save checkpoint. Note: Must guarantee that "save_ckpt_after_run" is registered in "after_run" hook. This method is called in: - ``auto_checkpoint`` (``torch_utils/checkpoint_helper.py``), which is designed for \ saving checkpoint whenever an exception raises. - ``serial_pipeline`` (``entry/serial_entry.py``). Used to save checkpoint when reaching \ new highest episode return. """ if ckpt_name is not None: self.ckpt_name = ckpt_name names = [h.name for h in self._hooks['after_run']] assert 'save_ckpt_after_run' in names idx = names.index('save_ckpt_after_run') self._hooks['after_run'][idx](self) self.ckpt_name = None
@property def learn_info(self) -> dict: """ Overview: Get current info dict, which will be sent to commander, e.g. replay buffer priority update, current iteration, hyper-parameter adjustment, whether task is finished, etc. Returns: - info (:obj:`dict`): Current learner info dict. """ ret = { 'learner_step': self._last_iter.val, 'priority_info': self.priority_info, 'learner_done': self._learner_done, } return ret @property def last_iter(self) -> CountVar: return self._last_iter @property def train_iter(self) -> int: return self._last_iter.val @property def monitor(self) -> 'TickMonitor': # noqa return self._monitor @property def log_buffer(self) -> dict: # LogDict return self._log_buffer @log_buffer.setter def log_buffer(self, _log_buffer: Dict[str, Dict[str, Any]]) -> None: self._log_buffer = _log_buffer @property def logger(self) -> logging.Logger: return self._logger @property def tb_logger(self) -> 'TensorBoradLogger': # noqa return self._tb_logger @property def exp_name(self) -> str: return self._exp_name @property def instance_name(self) -> str: return self._instance_name @property def rank(self) -> int: return self._rank @property def world_size(self) -> int: return self._world_size @property def policy(self) -> 'Policy': # noqa return self._policy @policy.setter def policy(self, _policy: 'Policy') -> None: # noqa """ Note: Policy variable monitor is set alongside with policy, because variables are determined by specific policy. """ self._policy = _policy if self._rank == 0: self._monitor = get_simple_monitor_type(self._policy.monitor_vars())(TickTime(), expire=10) if self._cfg.log_policy: self.info(self._policy.info()) @property def priority_info(self) -> dict: if not hasattr(self, '_priority_info'): self._priority_info = {} return self._priority_info @priority_info.setter def priority_info(self, _priority_info: dict) -> None: self._priority_info = _priority_info @property def ckpt_name(self) -> str: return self._ckpt_name @ckpt_name.setter def ckpt_name(self, _ckpt_name: str) -> None: self._ckpt_name = _ckpt_name
def create_learner(cfg: EasyDict, **kwargs) -> BaseLearner: """ Overview: Given the key(learner_name), create a new learner instance if in learner_mapping's values, or raise an KeyError. In other words, a derived learner must first register, then can call ``create_learner`` to get the instance. Arguments: - cfg (:obj:`EasyDict`): Learner config. Necessary keys: [learner.import_module, learner.learner_type]. Returns: - learner (:obj:`BaseLearner`): The created new learner, should be an instance of one of \ learner_mapping's values. """ import_module(cfg.get('import_names', [])) return LEARNER_REGISTRY.build(cfg.type, cfg=cfg, **kwargs) class TickMonitor(LoggedModel): """ Overview: TickMonitor is to monitor related info during training. Info includes: cur_lr, time(data, train, forward, backward), loss(total,...) These info variables are firstly recorded in ``log_buffer``, then in ``LearnerHook`` will vars in in this monitor be updated by``log_buffer``, finally printed to text logger and tensorboard logger. Interface: __init__, fixed_time, current_time, freeze, unfreeze, register_attribute_value, __getattr__ Property: time, expire """ data_time = LoggedValue(float) train_time = LoggedValue(float) total_collect_step = LoggedValue(float) total_step = LoggedValue(float) total_episode = LoggedValue(float) total_sample = LoggedValue(float) total_duration = LoggedValue(float) def __init__(self, time_: 'BaseTime', expire: Union[int, float]): # noqa LoggedModel.__init__(self, time_, expire) self.__register() def __register(self): def __avg_func(prop_name: str) -> float: records = self.range_values[prop_name]() _list = [_value for (_begin_time, _end_time), _value in records] return sum(_list) / len(_list) if len(_list) != 0 else 0 def __val_func(prop_name: str) -> float: records = self.range_values[prop_name]() return records[-1][1] for k in getattr(self, '_LoggedModel__properties'): self.register_attribute_value('avg', k, partial(__avg_func, prop_name=k)) self.register_attribute_value('val', k, partial(__val_func, prop_name=k)) def get_simple_monitor_type(properties: List[str] = []) -> TickMonitor: """ Overview: Besides basic training variables provided in ``TickMonitor``, many policies have their own customized ones to record and monitor. This function can return a customized tick monitor. Compared with ``TickMonitor``, ``SimpleTickMonitor`` can record extra ``properties`` passed in by a policy. Argumenst: - properties (:obj:`List[str]`): Customized properties to monitor. Returns: - simple_tick_monitor (:obj:`SimpleTickMonitor`): A simple customized tick monitor. """ if len(properties) == 0: return TickMonitor else: attrs = {} properties = [ 'data_time', 'train_time', 'sample_count', 'total_collect_step', 'total_step', 'total_sample', 'total_episode', 'total_duration' ] + properties for p_name in properties: attrs[p_name] = LoggedValue(float) return type('SimpleTickMonitor', (TickMonitor, ), attrs)