gt4sd.frameworks.gflownet.util module

Summary

Classes:

MPModelPlaceholder	This class can be used as a Model in a worker process, and translates calls to queries to the main process.
MPModelProxy	This class maintains a reference to an in-cuda-memory model, and creates a placeholder attribute which can be safely passed to multiprocessing DataLoader workers.

Functions:

wrap_model_mp

Construct a multiprocessing model proxy for torch DataLoaders so that only one process ends up making cuda calls and holding cuda tensors in memory.

Reference

class MPModelPlaceholder(in_queues, out_queues)

Bases: object

This class can be used as a Model in a worker process, and translates calls to queries to the main process.

__init__(in_queues, out_queues)

_check_init()

log_z(*a)

__call__(*a)

Call self as a function.

__dict__ = mappingproxy({'__module__': 'gt4sd.frameworks.gflownet.util', '__doc__': 'This class can be used as a Model in a worker process, and\n    translates calls to queries to the main process.\n    ', '__init__': <function MPModelPlaceholder.__init__>, '_check_init': <function MPModelPlaceholder._check_init>, 'log_z': <function MPModelPlaceholder.log_z>, '__call__': <function MPModelPlaceholder.__call__>, '__dict__': <attribute '__dict__' of 'MPModelPlaceholder' objects>, '__weakref__': <attribute '__weakref__' of 'MPModelPlaceholder' objects>, '__annotations__': {}})

__doc__ = 'This class can be used as a Model in a worker process, and\n    translates calls to queries to the main process.\n    '

__module__ = 'gt4sd.frameworks.gflownet.util'

__weakref__

list of weak references to the object (if defined)

class MPModelProxy(model, num_workers, cast_types)

Bases: object

This class maintains a reference to an in-cuda-memory model, and creates a placeholder attribute which can be safely passed to multiprocessing DataLoader workers.

This placeholder model sends messages accross multiprocessing queues, which are received by this proxy instance, which calls the model and sends the return value back to the worker.

Starts its own (daemon) thread. Always passes CPU tensors between processes.

__init__(model, num_workers, cast_types)

Construct a multiprocessing model proxy for torch DataLoaders.

Parameters

• model (Module) – a torch model which lives in the main process to which method calls are passed.

• num_workers (int) – number of workers.

• cast_types (tuple) – types that will be cast to cuda when received as arguments of method calls.

__del__()

run()

__dict__ = mappingproxy({'__module__': 'gt4sd.frameworks.gflownet.util', '__doc__': 'This class maintains a reference to an in-cuda-memory model, and\n    creates a `placeholder` attribute which can be safely passed to\n    multiprocessing DataLoader workers.\n\n    This placeholder model sends messages accross multiprocessing\n    queues, which are received by this proxy instance, which calls the\n    model and sends the return value back to the worker.\n\n    Starts its own (daemon) thread. Always passes CPU tensors between\n    processes.\n    ', '__init__': <function MPModelProxy.__init__>, '__del__': <function MPModelProxy.__del__>, 'run': <function MPModelProxy.run>, '__dict__': <attribute '__dict__' of 'MPModelProxy' objects>, '__weakref__': <attribute '__weakref__' of 'MPModelProxy' objects>, '__annotations__': {}})

__doc__ = 'This class maintains a reference to an in-cuda-memory model, and\n    creates a `placeholder` attribute which can be safely passed to\n    multiprocessing DataLoader workers.\n\n    This placeholder model sends messages accross multiprocessing\n    queues, which are received by this proxy instance, which calls the\n    model and sends the return value back to the worker.\n\n    Starts its own (daemon) thread. Always passes CPU tensors between\n    processes.\n    '

__module__ = 'gt4sd.frameworks.gflownet.util'

__weakref__

list of weak references to the object (if defined)

wrap_model_mp(model, num_workers, cast_types)

Construct a multiprocessing model proxy for torch DataLoaders so that only one process ends up making cuda calls and holding cuda tensors in memory.

Parameters

• model (Module) – a torch model which lives in the main process to which method calls are passed.

• num_workers (int) – number of DataLoader workers.

• cast_types (tuple) – types that will be cast to cuda when received as arguments of method calls. torch.Tensor is cast by default.

Returns

a placeholder model whose method calls route arguments to the main process.

Return type

placeholder