flambe

Subpackages

Package Contents

class flambe.Component(**kwargs)

Bases: flambe.compile.registrable.Registrable

Class which can be serialized to yaml and implements compile

IMPORTANT: ALWAYS inherit from Component BEFORE torch.nn.Module

Automatically registers subclasses via Registrable and facilitates immediate usage in YAML with tags. When loaded, subclasses’ initialization is delayed; kwargs are wrapped in a custom schema called Schema that can be easily initialized later.

_flambe_version = 0.0.0
_config_str

Represent object’s architecture as a YAML string

Includes the extensions relevant to the object as well; NOTE: currently this section may include a superset of the extensions actually needed, but this will be changed in a future release.

run(self)

Run a single computational step.

When used in an experiment, this computational step should be on the order of tens of seconds to about 10 minutes of work on your intended hardware; checkpoints will be performed in between calls to run, and resources or search algorithms will be updated. If you want to run everything all at once, make sure a single call to run does all the work and return False.

Returns:True if should continue running later i.e. more work to do
Return type:bool
metric(self)

Override this method to enable scheduling and searching.

Returns:The metric to compare different variants of your Component
Return type:float
register_attrs(self, *names: str)

Set attributes that should be included in state_dict

Equivalent to overriding obj._state and obj._load_state to save and load these attributes. Recommended usage: call inside __init__ at the end: self.register_attrs(attr1, attr2, …) Should ONLY be called on existing attributes.

Parameters:*names (str) – The names of the attributes to register
Raises:AttributeError – If self does not have existing attribute with that name
static _state_dict_hook(self, state_dict: State, prefix: str, local_metadata: Dict[str, Any])

Add metadata and recurse on Component children

This hook is used to integrate with the PyTorch state_dict mechanism; as either nn.Module.state_dict or Component.get_state recurse, this hook is responsible for adding Flambe specific metadata and recursing further on any Component children of self that are not also nn.Modules, as PyTorch will handle recursing to the latter.

Flambe specific metadata includes the class version specified in the Component._flambe_version class property, the name of the class, the source code, and the fact that this class is a Component and should correspond to a directory in our hiearchical save format

Finally, this hook calls a helper _state that users can implement to add custom state to a given class

Parameters:
  • state_dict (State) – The state_dict as defined by PyTorch; a flat dictionary with compound keys separated by ‘.’
  • prefix (str) – The current prefix for new compound keys that reflects the location of this instance in the object hierarchy being represented
  • local_metadata (Dict[str, Any]) – A subset of the metadata relevant just to this object and its children
Returns:

The modified state_dict
Return type:

type
Raises:

ExceptionName – Why the exception is raised.
_add_registered_attrs(self, state_dict: State, prefix: str)
_state(self, state_dict: State, prefix: str, local_metadata: Dict[str, Any])

Add custom state to state_dict

Parameters:
  • state_dict (State) – The state_dict as defined by PyTorch; a flat dictionary with compound keys separated by ‘.’
  • prefix (str) – The current prefix for new compound keys that reflects the location of this instance in the object hierarchy being represented
  • local_metadata (Dict[str, Any]) – A subset of the metadata relevant just to this object and its children
Returns:

The modified state_dict
Return type:

State
get_state(self, destination: Optional[State] = None, prefix: str = '', keep_vars: bool = False)

Extract PyTorch compatible state_dict

Adds Flambe specific properties to the state_dict, including special metadata (the class version, source code, and class name). By default, only includes state that PyTorch nn.Module includes (Parameters, Buffers, child Modules). Custom state can be added via the _state helper method which subclasses should override.

The metadata _flambe_directories indicates which objects are Components and should be a subdirectory in our hierarchical save format. This object will recurse on Component and nn.Module children, but NOT torch.optim.Optimizer subclasses, torch.optim.lr_scheduler._LRScheduler subclasses, or any other arbitrary python objects.

Parameters:
  • destination (Optional[State]) – The state_dict as defined by PyTorch; a flat dictionary with compound keys separated by ‘.’
  • prefix (str) – The current prefix for new compound keys that reflects the location of this instance in the object hierarchy being represented
  • keep_vars (bool) – Whether or not to keep Variables (only used by PyTorch) (the default is False).
Returns:

The state_dict object
Return type:

State
Raises:

ExceptionName – Why the exception is raised.
_load_state_dict_hook(self, state_dict: State, prefix: str, local_metadata: Dict[str, Any], strict: bool, missing_keys: List[Any], unexpected_keys: List[Any], error_msgs: List[Any])

Load flambe-specific state

Parameters:
  • state_dict (State) – The state_dict as defined by PyTorch; a flat dictionary with compound keys separated by ‘.’
  • prefix (str) – The current prefix for new compound keys that reflects the location of this instance in the object hierarchy being represented
  • local_metadata (Dict[str, Any]) – A subset of the metadata relevant just to this object and its children
  • strict (bool) – Whether missing or unexpected keys should be allowed; should always be False in Flambe
  • missing_keys (List[Any]) – Missing keys so far
  • unexpected_keys (List[Any]) – Unexpected keys so far
  • error_msgs (List[Any]) – Any error messages so far
Raises:

LoadError – If the state for some object does not have a matching major version number
_load_registered_attrs(self, state_dict: State, prefix: str)
_load_state(self, state_dict: State, prefix: str, local_metadata: Dict[str, Any], strict: bool, missing_keys: List[Any], unexpected_keys: List[Any], error_msgs: List[Any])

Load custom state (that was included via _state)

Subclasses should override this function to add custom state that isn’t normally included by PyTorch nn.Module

Parameters:
  • state_dict (State) – The state_dict as defined by PyTorch; a flat dictionary with compound keys separated by ‘.’
  • prefix (str) – The current prefix for new compound keys that reflects the location of this instance in the object hierarchy being represented
  • local_metadata (Dict[str, Any]) – A subset of the metadata relevant just to this object and its children
  • strict (bool) – Whether missing or unexpected keys should be allowed; should always be False in Flambe
  • missing_keys (List[Any]) – Missing keys so far
  • unexpected_keys (List[Any]) – Unexpected keys so far
  • error_msgs (List[Any]) – Any error messages so far
load_state(self, state_dict: State, strict: bool = False)

Load state_dict into self

Loads state produced by get_state into the current object, recursing on child Component and nn.Module objects

Parameters:
  • state_dict (State) – The state_dict as defined by PyTorch; a flat dictionary with compound keys separated by ‘.’
  • strict (bool) – Whether missing or unexpected keys should be allowed; should ALWAYS be False in Flambe (the default is False).
Raises:

LoadError – If the state for some object does not have a matching major version number
classmethod load_from_path(cls, path: str, map_location: Union[torch.device, str] = None, use_saved_config_defaults: bool = True, **kwargs: Any)
save(self, path: str, **kwargs: Any)
classmethod to_yaml(cls, representer: Any, node: Any, tag: str)
classmethod from_yaml(cls: Type[C], constructor: Any, node: Any, factory_name: str)
classmethod precompile(cls: Type[C], **kwargs: Any)

Change kwargs before compilation occurs.

This hook is called after links have been activated, but before calling the recursive initialization process on all other objects in kwargs. This is useful in a number of cases, for example, in Trainer, we compile several objects ahead of time and move them to the GPU before compiling the optimizer, because it needs to be initialized with the model parameters after they have been moved to GPU.

Parameters:
  • cls (Type[C]) – Class on which method is called
  • **kwargs (Any) – Current kwargs that will be compiled and used to initialize an instance of cls after this hook is called
aggregate_extensions_metadata(self)

Aggregate extensions used in object hierarchy

TODO: remove or combine with schema implementation in refactor

classmethod compile(cls: Type[C], _flambe_custom_factory_name: Optional[str] = None, _flambe_extensions: Optional[Dict[str, str]] = None, _flambe_stash: Optional[Dict[str, Any]] = None, **kwargs: Any)

Create instance of cls after recursively compiling kwargs

Similar to normal initialization, but recursively initializes any arguments that should be compiled and allows overriding arbitrarily deep kwargs before initializing if needed. Also activates any Link instances passed in as kwargs, and saves the original kwargs for dumping to yaml later.

Parameters:**kwargs (Any) – Keyword args that should be forwarded to the initialization function (a specified factory, or the normal __new__ and __init__ methods)
Returns:An instance of the class cls
Return type:C
class flambe.Schema(component_subclass: Type[C], _flambe_custom_factory_name: Optional[str] = None, **keywords: Any)

Bases: MutableMapping[str, Any]

Holds and recursively initializes Component’s with kwargs

Holds a Component subclass and keyword arguments to that class’s compile method. When an instance is called it will perform the recursive compilation process, turning the nested structure of Schema’s into initialized Component objects

Implements MutableMapping methods to facilitate inspection and updates to the keyword args. Implements dot-notation access to the keyword args as well.

Parameters:
  • component_subclass (Type[Component]) – Subclass of Component that will be compiled
  • **keywords (Any) – kwargs passed into the Schema’s compile method

Examples

Create a Schema from a Component subclass

>>> class Test(Component):
...     def __init__(self, x=2):
...         self.x = x
...
>>> tp = Schema(Test)
>>> t1 = tp()
>>> t2 = tp()
>>> assert t1 is t2  # the same Schema always gives you same obj
>>> tp = Schema(Test) # create a new Schema
>>> tp['x'] = 3
>>> t3 = tp()
>>> assert t1.x == 3  # dot notation works as well
component_subclass

Subclass of Schema that will be compiled

Type:Type[Component]
keywords

kwargs passed into the Schema’s compile method

Type:Dict[str, Any]
__call__(self, stash: Optional[Dict[str, Any]] = None, **keywords: Any)
add_extensions_metadata(self, extensions: Dict[str, str])

Add extensions used when loading this schema and children

Uses component_subclass.__module__ to filter for only the single relevant extension for this object; extensions relevant for children are saved only on those children schemas directly. Use aggregate_extensions_metadata to generate a dictionary of all extensions used in the object hierarchy.

aggregate_extensions_metadata(self)

Aggregate extensions used in object hierarchy

contains(self, schema: Schema, original_link: Link)
__setitem__(self, key: str, value: Any)
__getitem__(self, key: str)
__delitem__(self, key: str)
__iter__(self)
__len__(self)
__getattr__(self, item: str)
__setattr__(self, key: str, value: Any)
__repr__(self)

Identical to super (schema), but sorts keywords

classmethod to_yaml(cls, representer: Any, node: Any, tag: str = '')
static serialize(obj: Any)

Return dictionary representation of schema

Includes yaml as a string, and extensions

Parameters:obj (Any) – Should be schema or dict of schemas
Returns:dictionary containing yaml and extensions dictionary
Return type:Dict[str, Any]
static deserialize(data: Dict[str, Any])

Construct Schema from dict returned by Schema.serialize

Parameters:data (Dict[str, Any]) – dictionary returned by Schema.serialize
Returns:Schema or dict of schemas (depending on yaml in data)
Return type:Any
flambe.save(obj: Any, path: str, compress: bool = False, pickle_only: bool = False, overwrite: bool = False, pickle_module=dill, pickle_protocol=DEFAULT_PROTOCOL) → None

Save Component object to given path

See save_state_to_file for a more detailed explanation

Parameters:
  • obj (Component) – The component to save.
  • path (str) – Location to save the file / save directory to
  • compress (bool) – Whether to compress the save file / directory via tar + gz
  • pickle_only (bool) – Use given pickle_module instead of the hiearchical save format (the default is False).
  • overwrite (bool) – If true, overwrites the contents of the given path to create a directory at that location
  • pickle_module (type) – Pickle module that has load and dump methods; dump should accept a pickle_protocol parameter (the default is dill).
  • pickle_protocol (type) – Pickle protocol to use; see pickle for more details (the default is 2).
flambe.load(path: str, map_location=None, auto_install=False, pickle_module=dill, **pickle_load_args)

Load object with state from the given path

Loads a flambe object by using the saved config files, and then loads the saved state into said object. See load_state_from_file for details regarding how the state is loaded from the save file or directory.

Parameters:
  • path (str) – Path to the save file or directory
  • map_location (type) – Location (device) where items will be moved. ONLY used when the directory save format is used. See torch.load documentation for more details (the default is None).
  • auto_install (bool) – If True, automatically installs extensions as needed.
  • pickle_module (type) – Pickle module that has load and dump methods; dump should accept a pickle_protocol parameter (the default is dill).
  • **pickle_load_args (type) – Additional args that pickle_module should use to load; see torch.load documentation for more details
Returns:

object with both the architecture (config) and state that was saved to path
Return type:

Component
Raises:

LoadError – If a Component object is not loadable from the given path because extensions are not installed, or the config is empty, nonexistent, or otherwise invalid.
flambe.save_state_to_file(state: State, path: str, compress: bool = False, pickle_only: bool = False, overwrite: bool = False, pickle_module=dill, pickle_protocol=DEFAULT_PROTOCOL) → None

Save state to given path

By default the state will be saved in directory structure that mirrors the object hierarchy, so that you can later inspect the save file and load individual components more easily. If you would like to compress this directory structure using tar + gz, set compress to True. You can also use pickle to write a single output file, more similar to how PyTorch’s save function operates.

Parameters:
  • state (State) – The state_dict as defined by PyTorch; a flat dictionary with compound keys separated by ‘.’
  • path (str) – Location to save the file / save directory to; This should be a new non-existent path; if the path is an existing directory and it contains files an exception will be raised. This is because the path includes the final name of the save file (if using pickle or compress) or the final name of the save directory.
  • compress (bool) – Whether to compress the save file / directory via tar + gz
  • pickle_only (bool) – Use given pickle_module instead of the hiearchical save format (the default is False).
  • overwrite (bool) – If true, overwrites the contents of the given path to create a directory at that location
  • pickle_module (type) – Pickle module that has load and dump methods; dump should accpet a pickle_protocol parameter (the default is dill).
  • pickle_protocol (type) – Pickle protocol to use; see pickle for more details (the default is 2).
Raises:

ValueError – If the given path exists, is a directory, and already contains some files.
flambe.load_state_from_file(path: str, map_location=None, pickle_module=dill, **pickle_load_args) → State

Load state from the given path

Loads a flambe save directory, pickled save object, or a compressed version of one of these two formats (using tar + gz). Will automatically infer the type of save format and if the directory structure is used, the serialization protocol version as well.

Parameters:
  • path (str) – Path to the save file or directory
  • map_location (type) – Location (device) where items will be moved. ONLY used when the directory save format is used. See torch.load documentation for more details (the default is None).
  • pickle_module (type) – Pickle module that has load and dump methods; dump should accept a pickle_protocol parameter (the default is dill).
  • **pickle_load_args (type) – Additional args that pickle_module should use to load; see torch.load documentation for more details
Returns:

state_dict that can be loaded into a compatible Component
Return type:

State
flambe.log(tag: str, data: ValueT, global_step: int, walltime: Optional[float] = None) → None

Log data to tensorboard and console (convenience function)

Inspects type of data and uses the appropriate wrapper for tensorboard to consume the data. Supports floats (scalar), dictionary mapping tags to gloats (scalars), and strings (text).

Parameters:
  • tag (str) – Name of data, used as the tensorboard tag
  • data (ValueT) – The scalar or text to log
  • global_step (int) – Iteration number associated with data
  • walltime (Optional[float]) – Walltime for data (the default is None).

Examples

Normally you would have to do the following to log a scalar >>> import logging; from flambe.logging import ScalarT >>> logger = logging.getLogger(__name__) >>> logger.info(ScalarT(tag, data, step, walltime)) But this method allows you to write a more concise statement with a common interface >>> from flambe.logging import log >>> log(tag, data, step)

flambe.__version__

/$$$$$$$$ /$$ /$$ /$ | $$_____/| $$ | $$ /$ | $$ | $$ /$$$$$$ /$$$$$$/$$$$ | $$$$$$$ /$$$$$$ | $$$$$ | $$ |____ $$| $$_ $$_ $$| $$__ $$ /$$__ $$ | $$__/ | $$ /$$$$$$$| $$ $$ $$| $$ $$| $$$$$$$$ | $$ | $$ /$$__ $$| $$ | $$ | $$| $$ | $$| $$_____/ | $$ | $$| $$$$$$$| $$ | $$ | $$| $$$$$$$/| $$$$$$$ |__/ |__/ _______/|__/ |__/ |__/|_______/ _______/