## @package checkpoint # Module caffe2.python.checkpoint from __future__ import absolute_import from __future__ import division from __future__ import print_function from __future__ import unicode_literals import os import logging from caffe2.python import core, context from caffe2.python.net_builder import ops from caffe2.python.task import Node, Task, TaskGroup, TaskOutput, WorkspaceType logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) # The name of the special net that is used to store all the blob names in the # workspace. __BLOB_NAMES_NET__ = 'get_blob_list' @context.define_context() class Job(object): """ A Job defines three TaskGroups: the `init_group`, the `epoch_group` and the `exit_group` which will be run by a JobRunner. The `init_group` will be run only once at startup. Its role is to initialize globally persistent blobs such as model weights, accumulators and data file lists. The `epoch_group` will be run in a loop after init_group. The loop will exit when any of the stop signals added with `add_stop_signal` is True at the end of an epoch. The `exit_group` will be run only once at the very end of the job, when one of the stopping criterias for `epoch_group` was met. The role of this group is save the results of training in the end of the job. Jobs are context-driven, so that Tasks can be added to the active Job without having to explicitly pass the job object around. Example of usage: def build_reader(partitions): with Job.current().init_group: reader = HiveReader(init_reader, ..., partitions) Task(step=init_reader) with Job.current().epoch_group: limited_reader = ReaderWithLimit(reader, num_iter=10000) data_queue = pipe(limited_reader, num_threads=8) Job.current().add_stop_signal(limited_reader.data_finished()) return data_queue def build_hogwild_trainer(reader, model): with Job.current().init_group: Task(step=model.param_init_net) with Job.current().epoch_group: pipe(reader, processor=model, num_threads=8) with Job.current().exit_group: Task(step=model.save_model_net) with Job() as job: reader = build_reader(partitions) model = build_model(params) build_hogwild_trainer(reader, model) """ def __init__(self, init_group=None, epoch_group=None, exit_group=None, stop_signals=None, nodes_to_checkpoint=None): self.init_group = init_group or TaskGroup( workspace_type=WorkspaceType.GLOBAL) self.epoch_group = epoch_group or TaskGroup() self.exit_group = exit_group or TaskGroup() self.stop_signals = stop_signals or [] self._nodes_to_checkpoint = nodes_to_checkpoint def nodes_to_checkpoint(self): if self._nodes_to_checkpoint: return self._nodes_to_checkpoint else: return self.init_group.used_nodes() def compile(self, session_class): return Job( init_group=session_class.compile(self.init_group), epoch_group=session_class.compile(self.epoch_group), exit_group=session_class.compile(self.exit_group), stop_signals=self.stop_signals, nodes_to_checkpoint=self.nodes_to_checkpoint()) def __enter__(self): self.epoch_group.__enter__() return self def __exit__(self, *args): self.epoch_group.__exit__() def add_stop_signal(self, output): if isinstance(output, core.BlobReference): t = Task(outputs=[output], group=self.epoch_group) output = t.outputs()[0] assert isinstance(output, TaskOutput) self.stop_signals.append(output) class CheckpointManager(object): """ Controls saving and loading of workspaces on every epoch boundary of a job. If a CheckpointManager instance is passed to JobRunner, then JobRunner will call `init`, `read` and `save` at different moments in between epoch runs. """ def __init__(self, db, db_type): self._db = db self._db_type = db_type # make sure these blobs are the first in the checkpoint file. self._net = core.Net('!!checkpoint_mngr') self._blob_names = self._net.AddExternalInput('blob_names') self._names_output = None def init(self, nodes=None, retrieve_from_epoch=None): """ Build a Task that will be run once after the job's `init_group` is run. This task will determine which blobs need to be checkpointed. If retrieve_from_epoch is not None, then the checkpoint metadata is retrieved from a previously saved checkpoint. """ assert nodes is None or len(nodes) == 1, ( 'CheckpointManager only supports single node.') with Task(outputs=[self._blob_names]) as task: if retrieve_from_epoch is None: ops.GetAllBlobNames( [], self._blob_names, include_shared=False) else: ops.Load( [], self._blob_names, db=self._db_name(retrieve_from_epoch), db_type=self._db_type, absolute_path=True) self._names_output = task.outputs()[0] return task def blob_list(self): assert self._names_output return self._names_output.fetch().tolist() def _db_name(self, epoch): return '%s.%06d' % (self._db, epoch) def load(self, epoch): """ Build a Task that will be run by JobRunner when the job is to be resumed from a given epoch. This task will run a Load op that will load and deserialize all relevant blobs from a persistent storage. """ logger.info('Load from %s' % self._db_name(epoch)) with Task() as task: ops.Load( [], self.blob_list(), db=self._db_name(epoch), db_type=self._db_type, absolute_path=True) return task def load_blobs_from_checkpoint(self, blob_names, epoch): """ Builds a Task that loads only the necessary blobs from a checkpoint of the given epoch. The necessary blobs are given in the blob_names argument. Args: blob_names: A list of strings. Each string is the name of a blob. epoch: The checkpoint epoch to load from. Returns: A Task which loads the specified blobs from the checkpoint of the given epoch. """ logger.info('Load from %s' % self._db_name(epoch)) with Task() as task: ops.Load( [], blob_names, db=self._db_name(epoch), db_type=self._db_type, absolute_path=True, allow_incomplete=True) return task def check_db_exists(self, epoch): logger.info('Check existence of %s' % self._db_name(epoch)) with Task() as task: existence = ops.Const(False) ops.DBExists( [], [existence], db_name=self._db_name(epoch), db_type=self._db_type, absolute_path=True) task.add_output(existence) return task def save(self, epoch): """ Build a Task that is run once after `init_group` and after each epoch is run. This will execute a Save ops to serialize and persist blobs present in the global workspaace. """ logger.info('Save to %s' % self._db_name(epoch)) with Task() as task: ops.Save( self.blob_list(), [], db=self._db_name(epoch), db_type=self._db_type, absolute_path=True) return task class MultiNodeCheckpointManager(object): """ Coordinates checkpointing and checkpointing across multiple nodes. Each of `init`, `load` and `save` will build TaskGroups which will trigger checkpointing on each of the nodes involved in a distributed job. """ def __init__( self, db_prefix, db_type, node_manager_class=CheckpointManager): self._node_manager_class = node_manager_class self._node_managers = None self._db_prefix = db_prefix self._db_type = db_type def _task_group(self, func, *args, **kw): assert self._node_managers is not None, 'init must be called first.' with TaskGroup(WorkspaceType.GLOBAL) as task_group: for node, manager in self._node_managers: with Node(node): func(manager, *args, **kw) return task_group def init(self, nodes, retrieve_from_epoch=None): if self._node_managers is not None: assert [node for node, _ in self._node_managers] == nodes return self._node_managers = [] for node in nodes: with Node(node): manager = self._node_manager_class( db=os.path.join(self._db_prefix, node), db_type=self._db_type) self._node_managers.append((node, manager)) return self._task_group( self._node_manager_class.init, nodes=[node], retrieve_from_epoch=retrieve_from_epoch) def load(self, epoch): return self._task_group(self._node_manager_class.load, epoch) def load_blobs_locally(self, nodes, blob_names, epoch, session): """Loads the necessary blobs from the checkpoints to the current node. Args: blob_names: A list of strings. Each string is the name of a blob. epoch: An integer. The checkpoint epoch to load from. session: A Session object to execute the Load ops. """ if self._node_managers is not None: assert [node for node, _ in self._node_managers] == nodes else: self._node_managers = [] for node in nodes: with Node(node): manager = self._node_manager_class( db=os.path.join(self._db_prefix, node), db_type=self._db_type) self._node_managers.append((node, manager)) assert self._node_managers is not None, 'must initialize node managers' for _, manager in self._node_managers: existence_task = manager.check_db_exists(epoch) session.run(existence_task) existence = existence_task.outputs()[0].fetch() if not existence: logger.info('DB %s does not exist!' % manager._db_name(epoch)) return False load_task = manager.load_blobs_from_checkpoint(blob_names, epoch) session.run(load_task) logger.info('Successfully loaded from checkpoints.') return True def save(self, epoch): return self._task_group(self._node_manager_class.save, epoch) class JobRunner(object): """ Implement the runtime logic for jobs with checkpointing at the level of epoch. Can be used to run either single-host or distributed jobs. Job runner is a callable to be called once from the client, passing a Session as argument. This call will block until the Job execution is complete. If a checkpoint_manager is passed, checkpoints will be taken after initialization and after each epoch execution. If, in addition, `resume_from_epoch` is an epoch number, the corresponding checkpoint will be loaded and job execution will continue from the given epoch. In this case, the job's init_group will not be run. Refer to checkpoint_test.py for an example. """ def __init__(self, job, checkpoint_manager=None, resume_from_epoch=None): self.resume_from_epoch = resume_from_epoch self.checkpoint = checkpoint_manager self.job = job def __call__(self, client): from_scratch = self.resume_from_epoch is None if from_scratch: client.run(self.job.init_group) if self.checkpoint: logger.info('Preparing checkpoint ...') client.run(self.checkpoint.init( self.job.nodes_to_checkpoint(), retrieve_from_epoch=self.resume_from_epoch)) if from_scratch: logger.info('Saving first checkpoint ...') client.run(self.checkpoint.save(0)) logger.info('First checkpoint saved.') else: logger.info('Loading checkpoint for epoch {} ...'.format( self.resume_from_epoch)) client.run(self.checkpoint.load(self.resume_from_epoch)) logger.info('Checkpoint loaded.') epoch = 1 if from_scratch else self.resume_from_epoch + 1 while True: logger.info('Starting epoch %d.' % epoch) client.run(self.job.epoch_group) logger.info('Ran epoch %d.' % epoch) stop_signals = [o.fetch() for o in self.job.stop_signals] if self.checkpoint: logger.info('Saving checkpoint ...') client.run(self.checkpoint.save(epoch)) logger.info('Checkpoint saved.') if any(stop_signals): logger.info('Stopping.') break epoch += 1 client.run(self.job.exit_group) return epoch def load_blobs_from_checkpoints(self, blob_names, epoch, session): """Loads the necessary blobs from the checkpoints. Checkpoints store the snapshots of the workspace in each node. Sometimes we only need to load a subset of the blobs from the checkpoints. One common scenario is to load only the model blobs from the checkpoints for evaluation purpose. Given the names of the necessary blobs, this function goes over all the checkpoints of all the nodes, but only loads the blobs specified in the blob_names to the current workspace. Args: blob_names: A list of strings. Each string is the name of a blob. epoch: An integer. The checkpoint epoch to load from. session: A Session object to execute the load ops. Raises: ValueError: When the checkpoint manager is invalid. """ if not self.checkpoint: raise ValueError('Checkpoint manager is None') logger.info('Loading checkpoint for epoch {} ...'.format(epoch)) return self.checkpoint.load_blobs_locally(self.job.nodes_to_checkpoint(), blob_names, epoch, session) def epoch_limiter(num_epochs): """ Creates a task that will output True when a given number of epochs has finished. """ with Job.current().init_group: init_net = core.Net('epoch_counter_init') counter = init_net.CreateCounter([], init_count=num_epochs - 1) Task(step=init_net) epoch_net = core.Net('epoch_countdown') finished = epoch_net.CountDown(counter) output = Task(step=epoch_net, outputs=finished).outputs()[0] Job.current().add_stop_signal(output)