julius/prototorch
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Cleanup components

Browse Source
This commit is contained in:
Jensun Ravichandran 2021-05-17 16:58:57 +02:00
parent dc6248413c
commit 503ef0e05f
1 changed files with 16 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
prototorch/components/components.py
View File

@ -16,21 +16,23 @@ from torch.nn.parameter import Parameter
class Components(torch.nn.Module): class Components(torch.nn.Module):
"""Components is a set of learnable Tensors.""" """Components is a set of learnable Tensors."""
def __init__(self, def __init__(self,
number_of_components=None, ncomps=None,
initializer=None, initializer=None,
*, *,
initialized_components=None): initialized_components=None):
super().__init__() super().__init__()
self.ncomps = ncomps
# Ignore all initialization settings if initialized_components is given. # Ignore all initialization settings if initialized_components is given.
if initialized_components is not None: if initialized_components is not None:
self.register_parameter("_components", self.register_parameter("_components",
Parameter(initialized_components)) Parameter(initialized_components))
if number_of_components is not None or initializer is not None: if ncomps is not None or initializer is not None:
wmsg = "Arguments ignored while initializing Components" wmsg = "Arguments ignored while initializing Components"
warnings.warn(wmsg) warnings.warn(wmsg)
else: else:
self._initialize_components(number_of_components, initializer) self._initialize_components(initializer)
def _precheck_initializer(self, initializer): def _precheck_initializer(self, initializer):
if not isinstance(initializer, ComponentsInitializer): if not isinstance(initializer, ComponentsInitializer):
@ -39,9 +41,9 @@ class Components(torch.nn.Module):
f"You have provided: {initializer=} instead." f"You have provided: {initializer=} instead."
raise TypeError(emsg) raise TypeError(emsg)
def _initialize_components(self, number_of_components, initializer): def _initialize_components(self, initializer):
self._precheck_initializer(initializer) self._precheck_initializer(initializer)
_components = initializer.generate(number_of_components) _components = initializer.generate(self.ncomps)
self.register_parameter("_components", Parameter(_components)) self.register_parameter("_components", Parameter(_components))
@property @property
@ -72,17 +74,16 @@ class LabeledComponents(Components):
self._labels = component_labels self._labels = component_labels
else: else:
_labels = self._initialize_labels(distribution) _labels = self._initialize_labels(distribution)
super().__init__(number_of_components=len(_labels), super().__init__(len(_labels), initializer=initializer)
initializer=initializer) self.register_buffer("_labels", _labels)
self.register_buffer('_labels', _labels)
def _initialize_components(self, number_of_components, initializer): def _initialize_components(self, initializer):
if isinstance(initializer, ClassAwareInitializer): if isinstance(initializer, ClassAwareInitializer):
self._precheck_initializer(initializer) self._precheck_initializer(initializer)
self._components = Parameter( _components = initializer.generate(self.ncomps, self.distribution)
initializer.generate(number_of_components, self.distribution)) self.register_parameter("_components", Parameter(_components))
else: else:
super()._initialize_components(number_of_components, initializer) super()._initialize_components(initializer)
def _initialize_labels(self, distribution): def _initialize_labels(self, distribution):
if type(distribution) == tuple: if type(distribution) == tuple:
@ -130,14 +131,12 @@ class ReasoningComponents(Components):
self.register_parameter("_reasonings", reasonings) self.register_parameter("_reasonings", reasonings)
else: else:
self._initialize_reasonings(reasonings) self._initialize_reasonings(reasonings)
super().__init__(number_of_components=len(self._reasonings), super().__init__(len(self._reasonings), initializer=initializer)
initializer=initializer)
def _initialize_reasonings(self, reasonings): def _initialize_reasonings(self, reasonings):
if type(reasonings) == tuple: if type(reasonings) == tuple:
num_classes, number_of_components = reasonings nclasses, ncomps = reasonings
reasonings = ZeroReasoningsInitializer(num_classes, reasonings = ZeroReasoningsInitializer(nclasses, ncomps)
number_of_components)
_reasonings = reasonings.generate() _reasonings = reasonings.generate()
self.register_parameter("_reasonings", _reasonings) self.register_parameter("_reasonings", _reasonings)