def create_trainer( models: Union[torch.nn.Module, Mapping[str, torch.nn.Module]], optimizers: Union[torch.optim.Optimizer, Mapping[str, torch.optim.Optimizer]], max_epochs: int, *, extensions: Optional[Sequence[Union[ 'extension.ExtensionLike', 'extension.ExtensionEntry']]] = None, out_dir: str = 'result', stop_trigger: 'TriggerLike' = None, writer: Optional['writing.Writer'] = None, evaluator: Optional[Union['Evaluator', Tuple['Evaluator', 'TriggerLike'], Mapping[str, Union['Evaluator', Tuple[ 'Evaluator', 'TriggerLike']]]]] = None, device: 'DeviceLike' = 'cpu', logic: Optional[handler_module.BaseLogic] = None, transform_model: Callable[[str, torch.nn.Module], torch.nn.Module] = default_transform_model, handler_class: Optional[Type[handler_module.Handler]] = None, options: Optional[Dict[str, Any]] = None, runtime_options: Optional[Mapping[str, Any]] = None, profile: Optional[ torch.profiler.profile] = None, # type: ignore[name-defined] ) -> 'Trainer': """Creates a trainer object. Args: models: Map of string to Module or an actual Module. optimizers: Map of string to Optimizer or an actual Optimizer. max_epochs: Number of epochs in the whole training loop. Ignored if `stop_trigger` is passed as a kwarg. extensions: List of extensions to be registered to the trainer. out_dir: Output directory (default: ``result``). stop_trigger (trigger, optional): Trigger that can be consulted to determine wether training has concluded. The default is an interval trigger set to `max_epochs` writer: Writer that can be used by extensions to write data to custom filesystems. evaluator: Evaluator that is used in evaluation phase. If `None` is given, the evaluation is skipped. Evaluators can be created with :func:`pytorch_pfn_extras.engine.create_evaluator`. device (str or torch.device): Device name used for selecting a corresponding runtime class. logic: A logic object. If `None` is given, an logic object is instantiated from the default logic class. transform_model: A function to transform a model structure, often used to unwrap the a module from DDP module. handler_class: A handler class that instantiates a handler object. If `None` is given, `ppe.handler.Handler` is used as a default handler class. options: Options that are set to the handler and logic object. See the documentation of `ppe.handler.Handler` and `ppe.handler.Logic` for details. runtime_options: Options that are set to the runtime object. See the documentation of `ppe.runtime.PyTorchRuntime` for details. profile: A `torch.profiler.profile` object to collect the performance metrics. """ options = options.copy() if options else {} # TODO(kmaehashi): deprecate specifying 'runtime' key in options runtime_options = dict( runtime_options if runtime_options else options.pop('runtime', {})) logic = handler_module.Logic() if logic is None else logic handler_class = handler_class if handler_class else handler_module.Handler entry_runtime_cls = runtime_registry.get_runtime_class_for_device_spec( device) entry_runtime = entry_runtime_cls(device, runtime_options) handler = handler_class(logic, entry_runtime, {}) # Handler options are popped first, then pass the remainings to Logic. handler.consume_options(options) logic.consume_options(options) if len(options) > 0: raise ValueError('Unknown options: ', options) from pytorch_pfn_extras.training._trainer import Trainer return Trainer( handler, evaluator=evaluator, models=models, optimizers=optimizers, max_epochs=max_epochs, extensions=extensions, out_dir=out_dir, stop_trigger=stop_trigger, writer=writer, transform_model=transform_model, profile=profile, )
def create_evaluator( models: Union[torch.nn.Module, Mapping[str, torch.nn.Module]], *, progress_bar: bool = False, device: 'DeviceLike' = 'cpu', metrics: Optional[Sequence['MetricType']] = None, logic: Optional[handler_module.Logic] = None, handler_class: Optional[Type[handler_module.Handler]] = None, options: Optional[Dict[str, Any]] = None, runtime_options: Optional[Mapping[str, Any]] = None, profile: Optional[ torch.profiler.profile] = None, # type: ignore[name-defined] ) -> 'Evaluator': """Creates an evaluator object. The return value of this function is expected to be fed to `ppe.engine.create_trainer` as an argument. Args: models: Map of string to :class:`torch.nn.Module` or an actual Module. In most cases, this arugment is the same as the `model` arguemnt of `ppe.engine.create_trainer`. progress_bar: If `True`, a progress bar is enabled in evaluation. device (str or torch.device): Device name used for selecting a corresponding runtime class. metrics (list of metrics): List of metrics, which computes various quantities and update output for the reporting. logic: A logic object. If `None` is given, an logic object is instantiated from the default logic class. handler_class: A handler class that instantiates a handler object. If `None` is given, `ppe.handler.Handler` is used as a default handler class. options: Options that are set to the handler and logic object. See the documentation of `ppe.handler.Handler` and `ppe.handler.Logic` for details. runtime_options: Options that are set to the runtime object. See the documentation of `ppe.handler.Handler` for details. profile: A `torch.profiler.profile` object to collect the performance metrics. """ metrics = metrics if metrics else [] options = options.copy() if options else {} # TODO(kmaehashi): deprecate specifying 'runtime' key in options runtime_options = dict( runtime_options if runtime_options else options.pop('runtime', {})) logic = handler_module.Logic() if logic is None else logic handler_class = handler_class if handler_class else handler_module.Handler entry_runtime_cls = runtime_registry.get_runtime_class_for_device_spec( device) entry_runtime = entry_runtime_cls(device, runtime_options) handler = handler_class(logic, entry_runtime, options) # Handler options are popped first, then pass the remainings to Logic. handler.consume_options(options) logic.consume_options(options) if len(options) > 0: raise ValueError('Unknown options: ', options) from pytorch_pfn_extras.training._evaluator import Evaluator return Evaluator( handler, models=models, progress_bar=progress_bar, metrics=metrics, profile=profile, )