| # Copyright (c) 2016-present, Facebook, Inc. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| ############################################################################## |
| |
| ## @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) |
| |
| |
| @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) |
| |
| |
| def get_ckpt_filename(node_name, epoch): |
| """Returns the checkpoint filename. |
| |
| Args: |
| node_name: A string. The name of the node. |
| epoch: An integer. The checkpoint epoch. |
| |
| Returns: |
| ckpt_filename: A string. The filename of the checkpoint. |
| """ |
| return node_name + '.' + str(epoch) |
| |
| |
| def db_name(epoch, node_name, db_prefix, path_prefix=None): |
| """Returns the full db name where checkpoint files are saved. |
| |
| Args: |
| epoch: An integer. The checkpoint epoch. |
| node_name: A string. The name of the node. |
| db_prefix: A string. The prefix used to construct full db name. |
| path_prefix: A string. Optional param used to construct db name or path |
| where checkpoint files are are stored. |
| Returns: |
| db_name: A string. The absolute path of full_db_name where checkpoint |
| files are saved |
| """ |
| if path_prefix: |
| db_name = path_prefix + get_ckpt_filename(node_name, epoch) |
| else: |
| ckpt_filename = get_ckpt_filename(node_name, epoch) |
| db_name = os.path.join(db_prefix, ckpt_filename) |
| return db_name |
| |
| |
| 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. |
| |
| Args: |
| db_prefix: The prefix used to construct full db name. Since `absolute_path` |
| is set to True, this will be used as db_name in SaveOp. |
| node_name: Name of the node where this checkpoint_manager is used. |
| db_type: Type of database to use for storing checkpoint. |
| metadata_handler: An optional object capable of reading/writing |
| checkpoint info in storage of choice. |
| """ |
| def __init__(self, db_prefix, node_name, db_type, metadata_handler=None): |
| self._db_prefix = db_prefix |
| self._node_name = node_name |
| self._db_type = db_type |
| self._metadata_handler = metadata_handler |
| # 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 |
| |
| """ |
| Initialize the checkpoint manager. Determines all blobs that need to be saved |
| or loads from a checkpoint. |
| |
| Args: |
| nodes: An array of nodes where this checkpoint manager is running. Should |
| only contain a single node. |
| retrieve_from_epoch: Set to a number to load blobs from this epoch. |
| path_prefix: Used to construct db name or path where checkpoint files are |
| stored. |
| path_type: Indicate the type of path where checkpoint files are stored. |
| """ |
| def init( |
| self, |
| nodes=None, |
| retrieve_from_epoch=None, |
| path_prefix=None, |
| path_type=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.') |
| |
| self._path_prefix = path_prefix |
| self._path_type = path_type |
| |
| with Task(outputs=[self._blob_names]) as task: |
| if retrieve_from_epoch is None: |
| ops.GetAllBlobNames( |
| [], |
| self._blob_names, |
| include_shared=False) |
| else: |
| full_db_name = db_name(retrieve_from_epoch, |
| self._node_name, self._db_prefix, path_prefix) |
| db_type = path_type or self._db_type |
| logger.info("Initializing checkpoints from = %s" |
| % full_db_name) |
| ops.Load( |
| [], self._blob_names, |
| db=full_db_name, |
| db_type=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 load(self, epoch, path_prefix=None, path_type=None): |
| """ |
| 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. |
| """ |
| full_db_name = db_name(epoch, self._node_name, self._db_prefix, path_prefix) |
| db_type = path_type or self._db_type |
| logger.info("Loading checkpoints from = %s" % full_db_name) |
| with Task() as task: |
| ops.Load( |
| [], |
| self.blob_list(), |
| db=full_db_name, |
| db_type=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' % db_name(epoch, self._node_name, self._db_prefix)) |
| with Task() as task: |
| ops.Load( |
| [], |
| blob_names, |
| db=db_name(epoch, self._node_name, self._db_prefix), |
| 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' % |
| db_name(epoch, self._node_name, self._db_prefix)) |
| with Task() as task: |
| existence = ops.Const(False) |
| ops.DBExists( |
| [], |
| [existence], |
| db_name=db_name(epoch, self._node_name, self._db_prefix), |
| 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 workspace. |
| """ |
| logger.info('Saving to %s' % db_name(epoch, self._node_name, self._db_prefix)) |
| with Task() as task: |
| ops.Save( |
| self.blob_list(), [], |
| db=db_name(epoch, self._node_name, self._db_prefix), |
| db_type=self._db_type, absolute_path=True) |
| return task |
| |
| def write_checkpoint_metadata(self, epoch): |
| if self._metadata_handler is not None: |
| self._metadata_handler.write( |
| epoch=epoch, |
| db_type=self._db_type, |
| db_prefix=self._db_prefix, |
| path_type=self._path_type, |
| path_prefix=self._path_prefix, |
| node_names=[self._node_name], |
| ) |
| |
| def get_resume_from_epoch_id(self, user_epoch=None): |
| last_epoch = user_epoch |
| if self._metadata_handler is not None: |
| last_epoch = self._metadata_handler.last_epoch(user_epoch=user_epoch) |
| return last_epoch |
| |
| |
| 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. |
| |
| Args: |
| db_prefix: The prefix used to construct full db name. Since `absolute_path` |
| is set to True, this will be used as db_name in SaveOp. |
| db_type: Type of database to use for storing checkpoint. |
| metadata_handler: An optional object capable of reading/writing |
| checkpoint info in storage of choice. |
| """ |
| def __init__(self, db_prefix, db_type, metadata_handler=None): |
| self._node_managers = None |
| self._db_prefix = db_prefix |
| self._db_type = db_type |
| self._metadata_handler = metadata_handler |
| |
| 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 |
| |
| """ |
| Args: |
| nodes: An array of nodes where this checkpoint manager is running. |
| retrieve_from_epoch: Set to a number to load blobs from this epoch. |
| path_prefix: Used to construct db name or path where checkpoint files are |
| stored. |
| path_type: Indicate the type of path where checkpoint files are stored. |
| """ |
| def init( |
| self, nodes, retrieve_from_epoch=None, path_prefix=None, path_type=None |
| ): |
| if self._node_managers is not None: |
| assert [node for node, _ in self._node_managers] == nodes |
| return |
| self._node_managers = [] |
| self._path_prefix = path_prefix |
| self._path_type = path_type |
| self._node_names = [str(node) for node in nodes] |
| for node in nodes: |
| with Node(node): |
| manager = CheckpointManager( |
| db_prefix=self._db_prefix, |
| node_name=str(node), |
| db_type=self._db_type) |
| self._node_managers.append((node, manager)) |
| return self._task_group( |
| CheckpointManager.init, |
| nodes=[node], |
| retrieve_from_epoch=retrieve_from_epoch, |
| path_prefix=path_prefix, |
| path_type=path_type) |
| |
| def load(self, epoch, path_prefix=None, path_type=None): |
| return self._task_group( |
| CheckpointManager.load, |
| epoch, |
| path_prefix=path_prefix, |
| path_type=path_type) |
| |
| 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 = CheckpointManager( |
| db_prefix=self._db_prefix, |
| node_name=str(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!' % |
| db_name(epoch, manager._node_name, manager._db_prefix)) |
| 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 get_ckpt_db_name(self, node_name, epoch): |
| """Returns the DB name of the given node and the given epoch. |
| |
| The DB name is effectively the checkpoint path of the given node and |
| the given epoch. |
| |
| Args: |
| node_name: A string. The node name of interest. |
| epoch: An integer. The epoch of the checkpoint. |
| |
| Returns: |
| checkpoint_db_name: A string. The checkpoint path of the given |
| node and the given epoch. |
| """ |
| for node, manager in self._node_managers: |
| if str(node) == node_name: |
| return db_name(epoch, manager._node_name, manager._db_prefix) |
| |
| def save(self, epoch): |
| return self._task_group(CheckpointManager.save, epoch) |
| |
| def write_checkpoint_metadata(self, epoch): |
| if self._metadata_handler is not None: |
| self._metadata_handler.write( |
| epoch=epoch, |
| db_type=self._db_type, |
| db_prefix=self._db_prefix, |
| path_type=self._path_type, |
| path_prefix=self._path_prefix, |
| node_names=self._node_names, |
| ) |
| |
| def get_resume_from_epoch_id(self, user_epoch=None): |
| last_epoch = user_epoch |
| if self._metadata_handler is not None: |
| last_epoch = self._metadata_handler.last_epoch(user_epoch=user_epoch) |
| return last_epoch |
| |
| |
| class UploadTaskGroupBuilder(object): |
| """A simple class to upload checkpoints.""" |
| def build(self, epoch, checkpoint_manager): |
| """Builds the task group to upload checkpoints. |
| |
| Args: |
| epoch: An integer. The checkpoint epoch to be uploaded. |
| checkpoint_manager: Can be a CheckpointManager for single machine |
| or a MultiNodeCheckpointManager for multi-machine. The manager |
| that initializes/saves/loads checkpoints. |
| |
| Raises: |
| NotImplementedError: This base class only has the interface, |
| the implementation will be in the subclasses. |
| """ |
| raise NotImplementedError() |
| |
| |
| 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 master, passing a session |
| as an 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, |
| upload_task_group_builder=None): |
| """Initializes the JobRunner. |
| |
| Args: |
| job: A Job object. The job to be executed. |
| checkpoint_manager: Can be a CheckpointManager for single machine |
| or a MultiNodeCheckpointManager for multi-machine. The manager |
| that initializes/saves/loads checkpoints. |
| resume_from_epoch: An integer. The epoch to resume from. |
| upload_task_group_builder: A subclass of the |
| UploadTaskGroupBuilder. Creates a task group to upload |
| checkpoints. |
| """ |
| self.resume_from_epoch = resume_from_epoch |
| self.checkpoint_manager = checkpoint_manager |
| self.job = job |
| self.upload_task_group_builder = upload_task_group_builder |
| |
| def __call__(self, session): |
| """Runs the training flow. |
| |
| Args: |
| session: A Session object. Valid choises are: LocalSession, |
| LocalHostScheduler, and DistributedSession. It is used to |
| execute one TaskGroup a time. |
| """ |
| # identify the epoch we must resume from |
| if self.checkpoint_manager: |
| self.resume_from_epoch = self.checkpoint_manager.\ |
| get_resume_from_epoch_id(self.resume_from_epoch) |
| if self.resume_from_epoch is not None: |
| logger.info('Resuming from epoch {}'.format(self.resume_from_epoch)) |
| |
| # Initialize all the nodes. |
| from_scratch = self.resume_from_epoch is None |
| if from_scratch: |
| session.run(self.job.init_group) |
| |
| if self.checkpoint_manager: |
| logger.info('Preparing checkpoints ...') |
| session.run(self.checkpoint_manager.init( |
| self.job.nodes_to_checkpoint(), |
| retrieve_from_epoch=self.resume_from_epoch)) |
| # Save the first checkpoint before training starts, or resume from |
| # a previously saved checkpoint. |
| if from_scratch: |
| self.save_checkpoints(0, session) |
| else: |
| logger.info('Loading checkpoints for epoch {} ...'.format( |
| self.resume_from_epoch)) |
| session.run( |
| self.checkpoint_manager.load(self.resume_from_epoch)) |
| logger.info('Checkpoint loaded') |
| |
| logger.info("Finished initializing") |
| |
| # Start training. |
| epoch = 1 if from_scratch else self.resume_from_epoch + 1 |
| while True: |
| logger.info('Starting epoch %d' % epoch) |
| session.run(self.job.epoch_group) |
| logger.info('Finished epoch %d' % epoch) |
| stop_signals = [o.fetch() for o in self.job.stop_signals] |
| |
| if self.checkpoint_manager: |
| self.save_checkpoints(epoch, session) |
| |
| if any(stop_signals): |
| logger.info('Stopping') |
| break |
| epoch += 1 |
| logger.info('Finished training') |
| # Upload the checkpoints. |
| if (self.upload_task_group_builder): |
| upload_task_group = self.upload_task_group_builder.build( |
| epoch, self.checkpoint_manager) |
| session.run(upload_task_group) |
| logger.info('Finished uploading the checkpoints') |
| session.run(self.job.exit_group) |
| logger.info('Finished running the 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_manager: |
| raise ValueError('Checkpoint manager is None') |
| logger.info('Loading checkpoint for epoch {} ...'.format(epoch)) |
| return self.checkpoint_manager.load_blobs_locally( |
| self.job.nodes_to_checkpoint(), blob_names, epoch, session) |
| |
| def save_checkpoints(self, epoch, session): |
| """Triggers operation to save checkpoints |
| |
| This method will trigger the Save ops to serialize and persist the |
| blobs present in the global workspaace. |
| |
| Args: |
| epoch: An integer. The checkpoint epoch-id that we are saving. |
| session: A Session object to execute the save ops. |
| |
| Raises: |
| ValueError: When the checkpoint manager is invalid. |
| """ |
| if not self.checkpoint_manager: |
| raise ValueError('Checkpoint manager is None') |
| try: |
| logger.info('Saving checkpoints for epoch {}'.format(epoch)) |
| session.run(self.checkpoint_manager.save(epoch)) |
| self.checkpoint_manager.write_checkpoint_metadata(epoch) |
| logger.info('Checkpoints saved') |
| except Exception as ex: |
| logger.warning("Unable to write checkpoint for epoch {}. Error={}". |
| format(epoch, ex)) |
| |
| |
| 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) |
