pytorch_pfn_extras.training.extensions.snapshot_writers.ProcessQueueWriter#

class pytorch_pfn_extras.training.extensions.snapshot_writers.ProcessQueueWriter(savefun=<function save>, fs=None, out_dir='', task=None)#

Bases: QueueWriter[Process]

Snapshot writer that uses process queue.

This class creates a process and a queue by multiprocessing module. The process will be a consumer of this queue, and the main process will be a producer of this queue.

Note

Forking a new process from MPI process might be danger. Consider using ThreadQueueWriter instead of ProcessQueueWriter if you are using MPI.

See also

Methods

__init__([savefun, fs, out_dir, task])

consume(q)

create_consumer(q)

create_queue()

create_task(savefun)

finalize()

Finalizes the writer.

initialize(out_dir)

save(filename, out_dir, target, savefun, ...)

Parameters:

  • savefun (Callable[[...], None]) –

  • fs (Any) –

  • out_dir (str) –

  • task (Optional[Callable[[...], None]]) –

__init__(savefun=<function save>, fs=None, out_dir='', task=None)#
Parameters:

  • savefun (Callable[[...], None]) –

  • fs (Any) –

  • out_dir (str) –

  • task (Optional[Callable[[...], None]]) –

Return type:

None

create_consumer(q)#
Parameters:

q (Queue[Optional[Tuple[Callable[[...], None], str, str, Union[Sequence[Any], Mapping[str, Any]], Optional[Callable[[...], None]], bool]]]) –

Return type:

Process

create_queue()#
Return type:

Queue[Optional[Tuple[Callable[[…], None], str, str, Union[Sequence[Any], Mapping[str, Any]], Optional[Callable[[…], None]], bool]]]