pfnet / pytorch-pfn-extras

Supplementary components to accelerate research and development in PyTorch
https://medium.com/pytorch/migration-from-chainer-to-pytorch-8ed92c12c8
MIT License
272 stars 52 forks source link

Custom function extension #771

Open linshokaku opened 1 year ago

linshokaku commented 1 year ago

Implement a helper Extension that makes it easy to create a custom Extension that takes no manager as an argument and executes a pure custom function controlled by a trigger.

Custom functions that do not take manager as an argument can perform two main operations

As for side-effect operations, it is not possible to control them as a function of PPE due to the python specification, so no specific action is taken.

The report function would be more convenient if it could report on values returned by dict, following the PPE's handler module.

Therefore, I would implement it with the following interface

class UserFunctionExtension(Extension):
    priority = PRIORITY_WRITER
    def __init__(self, fn: Callable[..., Dict[str, Scalar]], args: Sequenc[Any], kwargs: Mapping[Any], report_prefix: str = None, priority: int = PRIORITY_WRITER) -> None:
        ...
linshokaku commented 1 year ago

https://github.com/pfnet/pytorch-pfn-extras/blob/c5b4d58c16667d60e980ca96c5839416d63d50f2/pytorch_pfn_extras/training/extension.py#L191

There is a need to clarify that it is easy to report() the results of the execution to avoid confusion caused by duplication of functionality with the conventional make_extension().