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

Device agnostic initialization of components.

Browse Source
This commit is contained in:
Alexander Engelsberger 2021-05-13 15:21:04 +02:00
parent aec422c277
commit 768d969f89
1 changed files with 18 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
prototorch/components/components.py
View File

@ -19,13 +19,13 @@ class Components(torch.nn.Module):
number_of_components=None, number_of_components=None,
initializer=None, initializer=None,
*, *,
initialized_components=None, initialized_components=None):
dtype=torch.float32):
super().__init__() super().__init__()
# 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._components = Parameter(initialized_components) self.register_parameter("_components",
Parameter(initialized_components))
if number_of_components is not None or initializer is not None: if number_of_components 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)
@ -41,13 +41,13 @@ class Components(torch.nn.Module):
def _initialize_components(self, number_of_components, initializer): def _initialize_components(self, number_of_components, initializer):
self._precheck_initializer(initializer) self._precheck_initializer(initializer)
self._components = Parameter( _components = initializer.generate(number_of_components)
initializer.generate(number_of_components)) self.register_parameter("_components", Parameter(_components))
@property @property
def components(self): def components(self):
"""Tensor containing the component tensors.""" """Tensor containing the component tensors."""
return self._components.detach().cpu() return self._components.detach()
def forward(self): def forward(self):
return self._components return self._components
@ -71,9 +71,10 @@ class LabeledComponents(Components):
super().__init__(initialized_components=components) super().__init__(initialized_components=components)
self._labels = component_labels self._labels = component_labels
else: else:
self._initialize_labels(distribution) _labels = self._initialize_labels(distribution)
super().__init__(number_of_components=len(self._labels), super().__init__(number_of_components=len(_labels),
initializer=initializer) initializer=initializer)
self.register_buffer('_labels', _labels)
def _initialize_components(self, number_of_components, initializer): def _initialize_components(self, number_of_components, initializer):
if isinstance(initializer, ClassAwareInitializer): if isinstance(initializer, ClassAwareInitializer):
@ -91,12 +92,12 @@ class LabeledComponents(Components):
labels = UnequalLabelsInitializer(distribution) labels = UnequalLabelsInitializer(distribution)
self.distribution = labels.distribution self.distribution = labels.distribution
self._labels = labels.generate() return labels.generate()
@property @property
def component_labels(self): def component_labels(self):
"""Tensor containing the component tensors.""" """Tensor containing the component tensors."""
return self._labels.detach().cpu() return self._labels.detach()
def forward(self): def forward(self):
return super().forward(), self._labels return super().forward(), self._labels
@ -123,8 +124,10 @@ class ReasoningComponents(Components):
*, *,
initialized_components=None): initialized_components=None):
if initialized_components is not None: if initialized_components is not None:
super().__init__(initialized_components=initialized_components[0]) components, reasonings = initialized_components
self._reasonings = initialized_components[1]
super().__init__(initialized_components=components)
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__(number_of_components=len(self._reasonings),
@ -136,7 +139,8 @@ class ReasoningComponents(Components):
reasonings = ZeroReasoningsInitializer(num_classes, reasonings = ZeroReasoningsInitializer(num_classes,
number_of_components) number_of_components)
self._reasonings = reasonings.generate() _reasonings = reasonings.generate()
self.register_parameter("_reasonings", _reasonings)
@property @property
def reasonings(self): def reasonings(self):
@ -145,7 +149,7 @@ class ReasoningComponents(Components):
Dimension NxCx2 Dimension NxCx2
""" """
return self._reasonings.detach().cpu() return self._reasonings.detach()
def forward(self): def forward(self):
return super().forward(), self._reasonings return super().forward(), self._reasonings