Merge branch 'borchify_namespace' into 'master' (a78e735a) · Commits · Desupervised / Borch

src/borch/init.py

+1 −1

Original line number	Diff line number	Diff line
		@@ -63,7 +63,7 @@ Examples:


		from borch.version import __version__
		from borch import infer, metrics, posterior
		from borch import infer, metrics, posterior, distributions
		from borch.random_variable import RandomVariable, RVPair, validate_args
		from borch.graph import Graph, as_tensor
		from borch.module import (

src/borch/module.py

+5 −6

Original line number	Diff line number	Diff line
		@@ -458,18 +458,17 @@ class Module(_Module):
		return out
		except AttributeError:
		posterior = self.__dict__.get("_modules", {}).get("posterior", None)
		prior = self.__dict__.get("_modules", {}).get("prior", None)
		if posterior is not None:
		param = getattr(self.posterior, name, None)
		if param is None and isinstance(
		getattr(self.prior, name, None), borch.RandomVariable
		getattr(prior, name, None), borch.RandomVariable
		):
		self.posterior.set_random_variable(
		name, getattr(self.prior, name), None
		)
		param = getattr(self.posterior, name)
		self.posterior.set_random_variable(name, getattr(prior, name), None)
		param = getattr(posterior, name)
		if isinstance(param, borch.RandomVariable):
		observed = self.observed.get(name)
		if self.observed.get(name) is not None:
		if observed is not None:
		param.tensor = observed
		self._used_rvs.add(name)
		# we convert it to just a tensor here to avoid potential bugs

src/borch/nn/init.py

+1 −1

Original line number	Diff line number	Diff line
		@@ -63,5 +63,5 @@ Notes:

		from borch.nn.torch_proxies import *
		from borch.nn import utils
		from borch.nn.borchify import borchify_module, borchify_network
		from borch.nn.borchify import borchify_module, borchify_network, borchify_namespace
		from borch.module import Module

src/borch/nn/borchify.py

+227 −38

Original line number	Diff line number	Diff line
		"""
		Functions to 'borchify' PyTorch modules/networks.
		"""
		# pylint: disable=too-many-arguments
		from typing import Optional
		from copy import deepcopy
		from functools import wraps

		from torch.nn import Module
		from torch import nn
		from torch.distributions import Distribution

		from borch.module import Module as BorchModule
		from borch.posterior import Posterior, Normal
		from borch.rv_factories import apply_rv_factory
		from borch.nn.torch_proxies import (
		get_rv_factory,
		_torch_proxy_class_factory,
		TORCH_BORCH_MAP,
		NOT_BORCHIFIED,
		)
		from borch.utils.module_utils import copy_module_attributes
		from borch.random_variable import RandomVariable
		from borch import posterior as borch_posterior
		from borch.module import Module as BorchModule, sample
		from borch.rv_factories import (
		apply_rv_factory,
		parameter_to_normal_rv,
		priors_to_rv,
		)
		from borch.utils.func_tools import assign_docs
		from borch.utils.namespace_tools import (
		create_augmented_classes,
		create_module,
		)


		def as_torch_module(cls, args, existing=None, *kwargs):
		"""Convert a module to a borch module"""
		obj = cls.__new__(cls, args, *kwargs)
		if existing is not None:
		copy_module_attributes(original=existing, new=obj)
		return obj
		def _get_priors_from_kwargs_(kwargs):
		priors = {}
		for key, val in kwargs.items():
		if isinstance(val, (Distribution, RandomVariable)):
		priors[key] = val
		for key in priors:
		kwargs.pop(key)
		return priors


		def default_rv_factory(_):
		"""Get the default random variable factory

		It creates a standard Gaussian.
		"""
		return parameter_to_normal_rv


		def borch_proxy_class(
		cls,
		get_rv_factory=default_rv_factory,
		doc_prefix="",
		borchify_submodules=False,
		get_extra_baseclasses=None,
		):
		"""
		Create a Bayesian version of a `torch.nn.Module`

		Args:
		cls: an uninstanciated class that is a subclass of `torch.nn.Module`
		get_rv_factory: function that takes a string as an argument and returns
		a function that creates random variables. See `borch.rv_factories`.
		doc_prefix (str): Extra documentation to append to the new class.
		"""
		if cls in BORCHIFY_REGISTRY:
		return BORCHIFY_REGISTRY[cls]
		extra_inheriors = (
		get_extra_baseclasses(cls) if get_extra_baseclasses is not None else []
		)

		@wraps(cls.__init__)
		def _init(self, args, posterior=None, *kwargs):
		# pylint: disable=redefined-outer-name
		if posterior is None:
		posterior = borch_posterior.Normal(loc_at_prior_mean=False)
		priors = _get_priors_from_kwargs_(kwargs)
		BorchModule.__init__(self, posterior=posterior)
		cls.__init__(self, args, *kwargs)
		self.posterior = posterior
		rv_factory = get_rv_factory(cls.__name__)
		for key in priors:
		if key not in self._parameters: # pylint: disable=protected-access
		raise ValueError(f"{key} does not match any parameters in the module")
		apply_rv_factory(
		self, lambda name, param: priors_to_rv(name, param, priors, rv_factory)
		)
		if borchify_submodules:
		_borchify_submodules_(self, rv_factory)
		sample(self, posterior=True, prior=True, redraw=False)

		new_cls = type(
		cls.__name__, (cls, BorchModule, *extra_inheriors), {"__init__": _init}
		)
		assign_docs(new_cls, cls, doc_prefix, False)
		return new_cls


		def borch_classes(
		module,
		get_rv_factory=default_rv_factory,
		doc_prefix="",
		ignore=None,
		parent=nn.Module,
		borchify_submodules=True,
		get_extra_baseclasses=None,
		):
		"""
		Create a Bayesian version of classes in a python module.

		Args:
		module: a python module that contains `torch.nn.Module`s you want a
		Bayesian version of.
		get_rv_factory: function that takes a string as an argument and returns
		a function that creates random variables. See `borch.rv_factories`.
		doc_prefix (str): Extra documentation to append to the new class.
		ignore (List(str)): List with string names of classes that should be skipped.
		parent (torch.nn.Module): a parent class you want to requite the subclasses
		to inherit from.
		borchify_submodules (bool): if submodules that are creates during the
		initialization method should also be borchified or not. Defaults to `False`.
		"""
		ignore = ignore if ignore is not None else []
		mappings = create_augmented_classes(
		module=module,
		parent=parent,
		class_factory=lambda cls: borch_proxy_class(
		cls,
		get_rv_factory=get_rv_factory,
		doc_prefix=doc_prefix,
		borchify_submodules=borchify_submodules,
		get_extra_baseclasses=get_extra_baseclasses,
		),
		ignore=ignore,
		)
		return mappings


		def borchify_namespace(
		module,
		get_rv_factory=default_rv_factory,
		doc_prefix="",
		ignore=None,
		parent=nn.Module,
		borchify_submodules=True,
		get_extra_baseclasses=None,
		):
		"""
		Create a new module that contains bayesian versions of the `torch.nn.Module`s.

		Args:
		module: a python module that contains `torch.nn.Module`s you want a
		Bayesian version of.
		get_rv_factory: function that takes a string as an argument and returns
		a function that creates random variables. See `borch.rv_factories`.
		doc_prefix (str): Extra documentation to append to the new class.
		ignore (List(str)): List with string names of classes that should be skipped.
		parent (torch.nn.Module): a parent class you want to requite the subclasses
		to inherit from.
		borchify_submodules (bool): if submodules that are creates during the
		initialization method should also be borchified or not. Defaults to `False`.

		Returns:
		A new python module where the `torch.nn.Modules` now also inherit from
		`borch.Module`.

		def as_borch_module(cls, args, existing=None, *kwargs):
		"""Convert a module to a borch module"""
		Examples:
		>>> import torch
		>>> bnn = borchify_namespace(torch.nn)
		>>> blinear = bnn.Linear(1,2)
		>>> type(blinear)
		<class 'borch.nn.borchify.Linear'>
		"""
		mappings = borch_classes(
		module=module,
		get_rv_factory=get_rv_factory,
		doc_prefix=doc_prefix,
		ignore=ignore,
		parent=parent,
		borchify_submodules=borchify_submodules,
		get_extra_baseclasses=get_extra_baseclasses,
		)
		classes = {mapping.name: mapping.augmented for mapping in mappings}
		namespace = dict(module.__dict__)
		namespace.update(classes)
		return create_module(module.__name__, module.__doc__, namespace)


		def _instanciate(cls, args, existing=None, *kwargs):
		"""Instantiate a class without calling the init method"""
		obj = cls.__new__(cls, args, *kwargs)
		if issubclass(cls, BorchModule):
		BorchModule.__init__(obj, args, *kwargs)
		if existing is not None:
		copy_module_attributes(original=existing, new=obj)
		@@ -41,15 +200,29 @@ def _register_rvs_with_posterior(module):
		getattr(module, key, None)


		_BORCHIFIED_CLASS_MAP = {}
		class Registry:
		"""Registry with mappings between different objects"""

		def __init__(self):
		self.registry = {}

		def register(self, kwargs, overwrite=False):
		"""Register a new mapping"""
		if isinstance(kwargs, list):
		kwargs = {mapping.original: mapping.augmented for mapping in kwargs}
		overlap = set(self.registry).intersection(set(kwargs))
		if overlap and not overwrite:
		raise RuntimeError(f"{overlap} is already register")
		self.registry.update(kwargs)

		def __getitem__(self, key):
		return self.registry.get(key)

		def _get_borchified_class(cls):
		if cls in TORCH_BORCH_MAP:
		return TORCH_BORCH_MAP[cls]
		if cls not in _BORCHIFIED_CLASS_MAP:
		_BORCHIFIED_CLASS_MAP[cls] = _torch_proxy_class_factory(cls)
		return _BORCHIFIED_CLASS_MAP[cls]
		def __contains__(self, key):
		return key in self.registry


		BORCHIFY_REGISTRY = Registry()


		def borchify_module(
		@@ -61,7 +234,7 @@ def borchify_module(
		Args:
		module: The ``Module`` object to be 'borchified'.
		rv_factory: A callable which, when passed a ``Parameter``, returns a
		``RandomVariable``, if None the default of ``borch.nn.torch_proxies``
		``RandomVariable``, if None the default of ``borch.nn.borchify``
		will be used.
		posterior: A posterior for which the borchified module should use. The default
		is ``Normal`` (see ``borch.posterior``).
		@@ -74,7 +247,7 @@ def borchify_module(
		>>> linear = torch.nn.Linear(3, 3) # create a linear module
		>>> blinear = borchify_module(linear)
		>>> type(blinear)
		<class 'borch.nn.torch_proxies.Linear'>
		<class 'borch.nn.borchify.Linear'>
		"""
		# pylint: disable=unexpected-keyword-arg

		@@ -86,12 +259,14 @@ def borchify_module(
		_register_rvs_with_posterior(new)
		return new

		cls_name = type(module).__name__
		new_module = _get_borchified_class(type(module))
		if new_module in NOT_BORCHIFIED:
		return as_torch_module(new_module, existing=module)
		new = as_borch_module(new_module, existing=module, posterior=posterior)
		temp_rv_factory = get_rv_factory(cls_name) if rv_factory is None else rv_factory
		cls_type = type(module)
		new_module = borch_proxy_class(cls_type)
		if not issubclass(new_module, BorchModule):
		return _instanciate(new_module, existing=module)
		new = _instanciate(new_module, existing=module, posterior=posterior)
		temp_rv_factory = (
		default_rv_factory(cls_type.__name__) if rv_factory is None else rv_factory
		)
		apply_rv_factory(new, temp_rv_factory)
		_register_rvs_with_posterior(new)
		return new
		@@ -109,7 +284,7 @@ def borchify_network(
		Args:
		module: The network to be borchified.
		rv_factory: A callable which, when passed a ``Parameter``, returns a
		``RandomVariable``, if None the default of ``borch.nn.torch_proxies``
		``RandomVariable``, if None the default of ``borch.nn.borchify``
		will be used.
		posterior_creator: A callable which creates a posterior. This will be used to
		create a new posterior for each module in the network.
		@@ -145,11 +320,10 @@ def borchify_network(
		>>> net = Net()
		>>> bnet = borchify_network(net)
		>>> type(bnet)
		<class 'borch.nn.torch_proxies.Net'>
		<class 'borch.nn.borchify.Net'>
		>>> type(bnet.linear)
		<class 'borch.nn.torch_proxies.Linear'>
		<class 'borch.nn.borchify.Linear'>
		"""
		# pylint: disable=protected-access
		if cache is None:
		cache = {}

		@@ -161,12 +335,27 @@ def borchify_network(
		posterior_creator = lambda: Normal(-4, loc_at_prior_mean=False)
		new = borchify_module(module, rv_factory, posterior_creator())
		cache[_id] = new
		_borchify_submodules_(
		module,
		rv_factory=rv_factory,
		posterior_creator=posterior_creator,
		target=new,
		cache=cache,
		)
		return new


		for name, mod in module._modules.items():
		def _borchify_submodules_(
		module, rv_factory=None, posterior_creator=None, target=None, cache=None
		):
		if target is None:
		target = module
		if cache is None:
		cache = {}
		for name, mod in module._modules.items(): # pylint: disable=protected-access
		if isinstance(module, BorchModule) and mod in module.internal_modules:
		continue
		adding = cache.get(
		id(mod), borchify_network(mod, rv_factory, posterior_creator, cache)
		)
		new.add_module(name, adding)
		return new
		target.add_module(name, adding)

src/borch/nn/torch_proxies.py

+47 −54

Original line number	Diff line number	Diff line
		@@ -35,23 +35,16 @@ Examples:
		"""

		# pylint: disable=protected-access,undefined-all-variable
		from functools import wraps
		import sys

		from torch import nn
		from torch.distributions import Distribution

		from borch.random_variable import RandomVariable
		from borch import posterior as borch_posterior
		from borch.module import Module as BorchModule, sample
		from borch.rv_factories import (
		apply_rv_factory,
		kaiming_normal_rv,
		parameter_to_normal_rv,
		priors_to_rv,
		)
		from borch.utils.func_tools import assign_docs
		from borch.utils.namespace_tools import create_augmented_classes
		from borch.utils.namespace_tools import extend_module
		from borch.nn.borchify import borch_classes, BORCHIFY_REGISTRY

		DOC_PREFIX = """This is a ppl class. Please see ``help(borch.nn)`` for more information.
		If one gives distribution as kwargs, where names match the parameters of the Module, they
		@@ -88,44 +81,6 @@ def get_rv_factory(cls_name):
		return RV_FACTORIES.get(cls_name, parameter_to_normal_rv)


		def _get_priors_from_kwargs_(kwargs):
		priors = {}
		for key, val in kwargs.items():
		if isinstance(val, (Distribution, RandomVariable)):
		priors[key] = val
		for key in priors:
		kwargs.pop(key)
		return priors


		def _torch_proxy_class_factory(cls):
		@wraps(cls.__init__)
		def _init(self, args, posterior=None, *kwargs):
		# pylint: disable=redefined-outer-name
		# This will become the new init function of borch.nn.<cls.__name__>
		if posterior is None:
		posterior = borch_posterior.Normal(loc_at_prior_mean=False)
		priors = _get_priors_from_kwargs_(kwargs)
		BorchModule.__init__(self, posterior=posterior)
		cls.__init__(self, args, *kwargs)
		self.posterior = posterior
		params = self._parameters
		rv_factory = get_rv_factory(cls.__name__)
		for key in priors:
		if key not in params:
		raise ValueError(f"{key} does not match any parameters in the module")
		apply_rv_factory(
		self, lambda name, param: priors_to_rv(name, param, priors, rv_factory)
		)
		# TODO do we want to set the values for the posterior samples??
		# if not there will be a lot of size mismatch when loading a state dict
		sample(self, posterior=True, prior=True, redraw=False)

		new_cls = type(cls.__name__, (cls, BorchModule), {"__init__": _init})
		assign_docs(new_cls, cls, DOC_PREFIX)
		return new_cls


		_NO_WEIGHTS_MODULE_NAMES = [
		"AlphaDropout",
		"BCELoss",
		@@ -167,6 +122,7 @@ _NO_WEIGHTS_MODULE_NAMES = [
		"LeakyReLU",
		"LogSigmoid",
		"LogSoftmax",
		"Mish",
		"MSELoss",
		"MarginRankingLoss",
		"ModuleDict",
		@@ -216,26 +172,62 @@ _NO_WEIGHTS_MODULE_NAMES = [
		"UpsamplingBilinear2d",
		"UpsamplingNearest2d",
		"ZeroPad2d",
		"HuberLoss",
		"LazyBatchNorm1d",
		"LazyBatchNorm2d",
		"LazyBatchNorm3d",
		"LazyInstanceNorm1d",
		"LazyInstanceNorm2d",
		"LazyInstanceNorm3d",
		"ReflectionPad3d",
		]


		class _RNNFlatWeights:
		"""Make sure we use getattr for the parameters"""

		@property
		def _flat_weights(self):
		return [
		(lambda wn: getattr(self, wn) if hasattr(self, wn) else None)(wn)
		for wn in self._flat_weights_names
		]

		@_flat_weights.setter
		def _flat_weights(self, val):
		pass


		def get_extra_baseclasses(cls):
		"""Extra base classes for torch proxies"""
		if issubclass(cls, nn.RNNBase):
		return [_RNNFlatWeights]
		return []

		TORCH_BORCH_MAP = create_augmented_classes(
		caller=__name__,
		module=nn,
		parent=nn.Module,
		class_factory=_torch_proxy_class_factory,
		ignore=_NO_WEIGHTS_MODULE_NAMES + ["Transformer"],

		_MAPPINGS = borch_classes(
		nn,
		get_rv_factory=get_rv_factory,
		doc_prefix=DOC_PREFIX,
		ignore=_NO_WEIGHTS_MODULE_NAMES,
		borchify_submodules=True,
		get_extra_baseclasses=get_extra_baseclasses,
		)
		TORCH_BORCH_MAP = {mapping.original: mapping.augmented for mapping in _MAPPINGS}
		BORCHIFY_REGISTRY.register(_MAPPINGS)
		extend_module(__name__, _MAPPINGS)


		NOT_BORCHIFIED = [
		getattr(nn, name) for name in _NO_WEIGHTS_MODULE_NAMES if hasattr(nn, name)
		]
		BORCHIFY_REGISTRY.register({v: v for v in NOT_BORCHIFIED})


		for _cls in NOT_BORCHIFIED:
		TORCH_BORCH_MAP[_cls] = _cls
		setattr(sys.modules[__name__], _cls.__name__, _cls)


		__all__ = [
		"AdaptiveAvgPool1d",
		"AdaptiveAvgPool2d",
		@@ -368,6 +360,7 @@ __all__ = [
		"Tanh",
		"Tanhshrink",
		"Threshold",
		"Transformer",
		"TransformerDecoder",
		"TransformerDecoderLayer",
		"TransformerEncoder",