flambe.compile

Package Contents

exception flambe.compile.RegistrationError[source]

Bases: Exception

Error thrown when acessing yaml tag on a non-registered class

Thrown when trying to access the default yaml tag for a class typically occurs when called on an abstract class

class flambe.compile.Registrable[source]

Bases: abc.ABC

Subclasses automatically registered as yaml tags

Automatically registers subclasses with the yaml loader by adding a constructor and representer which can be overridden

_yaml_tags :Dict[Any, List[str]]
_yaml_tag_namespace :Dict[Type, str]
_yaml_registered_factories :Set[str]
classmethod __init_subclass__(cls: Type[R], should_register: Optional[bool] = True, tag_override: Optional[str] = None, tag_namespace: Optional[str] = None, **kwargs: Mapping[str, Any])
static register_tag(class_: RT, tag: str, tag_namespace: Optional[str] = None)
static get_default_tag(class_: RT, factory_name: Optional[str] = None)

Retrieve default yaml tag for class cls

Retrieve the default tag (aka the last one, which will be the only one, or the alias if it exists) for use in yaml representation

classmethod to_yaml(cls, representer: Any, node: Any, tag: str)

Use representer to create yaml representation of node

See Component class, and experiment/options for examples

classmethod from_yaml(cls, constructor: Any, node: Any, factory_name: str)

Use constructor to create an instance of cls

See Component class, and experiment/options for examples

flambe.compile.alias(tag: str, tag_namespace: Optional[str] = None) → Callable[[RT], RT][source]

Decorate a Registrable subclass with a new tag

Can be added multiple times to give a class multiple aliases, however the top most alias tag will be the default tag which means it will be used when representing the class in YAML

flambe.compile.yaml[source]
flambe.compile.register(cls: Type[A], tag: str) → Type[A][source]

Safely register a new tag for a class

Similar to alias, but it’s intended to be used on classes that are not already subclasses of Registrable, and it is NOT a decorator

class flambe.compile.registrable_factory(fn: Any)[source]

Decorate Registrable factory method for use in the config

This Descriptor class will set properties that allow the factory method to be specified directly in the config as a suffix to the tag; for example:

class MyModel(Component):

    @registrable_factory
    def from_file(cls, path):
        # load instance from path
        ...
        return instance

defines the factory, which can then be used in yaml:

model: !MyModel.from_file
    path: some/path/to/file.pt
__set_name__(self, owner: type, name: str)
class flambe.compile.registration_context(namespace: str)[source]
__enter__(self)
__exit__(self, *args: Any)
__call__(self, func: Callable[..., Any])
class flambe.compile.MappedRegistrable[source]

Bases: flambe.compile.registrable.Registrable

classmethod to_yaml(cls, representer: Any, node: Any, tag: str)

Use representer to create yaml representation of node

classmethod from_yaml(cls, constructor: Any, node: Any, factory_name: str)

Use constructor to create an instance of cls

class flambe.compile.Schema(component_subclass: Type[C], _flambe_custom_factory_name: Optional[str] = None, **keywords: Any)[source]

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
class flambe.compile.Component(**kwargs)[source]

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

Bases: flambe.compile.registrable.Registrable

Delayed access to another object in an object hiearchy

Currently only supported in the context of Experiment but this may be updated in a future release

A Link delays the access of some property, or the calling of some method, until the Link is called. Links can be passed directly into a Component subclass compile, Component’s method called compile will automatically record the links and call them to access their values before running __new__ and __init__. The recorded links will show up in the config if yaml.dump() is called on your object hierarchy. This typically happens when logging individual configs during a grid search, and when serializing between multiple processes.

For example, if the schematic path is [‘model’, ‘encoder’] and the attribute path is [‘rnn’, ‘hidden_size’] then before the link can be compiled, the target attribute should be set to point to the model schema (this is handled automatically by Experiment) then, during compilation the child schema ‘encoder’ will be accessed, and finally the attribute encoder.rnn.hidden_size will be returned

Parameters:
  • schematic_path (Sequence[str]) – Path to the relevant schema denoted by dictionary-like bracket access e.g. [‘model’, ‘encoder’]
  • attr_path (Sequence[str]) – Path to the relevant attribute on the given schema (after it’s been compiled) using standard attribute dot notation e.g. [‘rnn’, ‘hidden_size’]
  • target (Optional[Schema]) – The root object corresponding to the first element in the schematic path; needs to be passed in here or set later before link can be resolved
  • local (bool) – if true, changes tune convert behavior to insert a dummy link; used for links to global variables (“resources” in config) (defaults to True)
root_schema :str
__repr__(self)
__call__(self)
classmethod to_yaml(cls, representer: Any, node: Any, tag: str)

Build contextualized link based on the root node

If the link refers to something inside of the current object hierarchy (as determined by the schema of _link_root_obj) then it will be represented as a link; if the link refers to something out-of-scope, i.e. not inside the current object hiearchy, then replace the link with the resolved value. If the value cannot be represented, pickle it and include a reference to its id in the object stash that will be saved alongside the config

classmethod from_yaml(cls, constructor: Any, node: Any, factory_name: str)
convert(self)
flambe.compile.dynamic_component(class_: Type[A], tag: str, tag_namespace: Optional[str] = None, parent_component_class: Type[Component] = Component) → Type[Component][source]

Decorate given class, creating a dynamic Component

Creates a dynamic subclass of class_ that inherits from Component so it will be registered with the yaml loader and receive the appropriate functionality (from_yaml, to_yaml and compile). class_ should not implement any of the aforementioned functions.

Parameters:
  • class (Type[A]) – Class to register with yaml and the compilation system
  • tag (str) – Tag that will be used with yaml
  • tag_namespace (str) – Namespace aka the prefix, used. e.g. for !torch.Adam torch is the namespace
Returns:

New subclass of _class and Component

Return type:

Type[Component]

flambe.compile.make_component(class_: type, tag_namespace: Optional[str] = None, only_module: Optional[str] = None, parent_component_class: Optional[Type] = None, exclude: Optional[List[str]] = None) → None[source]

Make class and all its children a Component

For example a call to make_component(torch.optim.Adam, “torch”) will make the tag !torch.Adam accessible in any yaml configs. This does NOT monkey patch (aka swizzle) torch, but instead creates a dynamic subclass which will be used by the yaml constructor i.e. only classes loaded from the config will be affected, anything imported and used in code.

Parameters:
  • class (type) – To be registered with yaml as a component, along with all its children
  • tag_prefix (str) – Added to beginning of all the tags
  • only_module (str) – Module prefix used to limit the scope of what gets registered
  • parent_component_class (Type) – Parent class to use for creating a new component class; should be a subclass of :class:~flambe.compile.Component (defaults to Component)
  • exclude (List[str], optional) – A list of modules to ignore
Returns:

Return type:

None

flambe.compile.all_subclasses(class_: Type[Any]) → Set[Type[Any]][source]

Return a set of all subclasses for a given class object

Recursively collects all subclasses of class_ down the object hierarchy into one set.

Parameters:class (Type[Any]) – Class to retrieve all subclasses for
Returns:All subclasses of class_
Return type:Set[Type[Any]]
flambe.compile.save(obj: Any, path: str, compress: bool = False, pickle_only: bool = False, overwrite: bool = False, pickle_module=dill, pickle_protocol=DEFAULT_PROTOCOL) → None[source]

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.compile.load(path: str, map_location=None, auto_install=False, pickle_module=dill, **pickle_load_args)[source]

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.compile.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[source]

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.compile.load_state_from_file(path: str, map_location=None, pickle_module=dill, **pickle_load_args) → State[source]

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

class flambe.compile.State[source]

Bases: collections.OrderedDict

A state object for Flambe.

_metadata :Dict[str, Any]