[FEATURE] Allow initialized_components to be a dataset

This commit is contained in:
Jensun Ravichandran 2021-06-04 22:13:36 +02:00
parent 729b20e9ab
commit 8200e1d3d8

View File

@ -11,6 +11,8 @@ from prototorch.components.initializers import (ClassAwareInitializer,
ZeroReasoningsInitializer) ZeroReasoningsInitializer)
from torch.nn.parameter import Parameter from torch.nn.parameter import Parameter
from .initializers import parse_data_arg
def get_labels_object(distribution): def get_labels_object(distribution):
if isinstance(distribution, dict): if isinstance(distribution, dict):
@ -56,24 +58,23 @@ class Components(torch.nn.Module):
wmsg = "Arguments ignored while initializing Components" wmsg = "Arguments ignored while initializing Components"
warnings.warn(wmsg) warnings.warn(wmsg)
else: else:
self._initialize_components(initializer, num_components) self._initialize_components(num_components, initializer)
@property @property
def num_components(self): def num_components(self):
# return len(self._components) return len(self._components)
return self._components.shape[0]
def _register_components(self, components): def _register_components(self, components):
self.register_parameter("_components", Parameter(components)) self.register_parameter("_components", Parameter(components))
def _initialize_components(self, initializer, num_components): def _initialize_components(self, num_components, initializer):
_precheck_initializer(initializer) _precheck_initializer(initializer)
_components = initializer.generate(num_components) _components = initializer.generate(num_components)
self._register_components(_components) self._register_components(_components)
def add_components(self, def add_components(self,
initializer=None,
num=1, num=1,
initializer=None,
*, *,
initialized_components=None): initialized_components=None):
if initialized_components is not None: if initialized_components is not None:
@ -114,7 +115,8 @@ class LabeledComponents(Components):
*, *,
initialized_components=None): initialized_components=None):
if initialized_components is not None: if initialized_components is not None:
components, component_labels = initialized_components components, component_labels = parse_data_arg(
initialized_components)
super().__init__(initialized_components=components) super().__init__(initialized_components=components)
self._labels = component_labels self._labels = component_labels
else: else:
@ -134,16 +136,16 @@ class LabeledComponents(Components):
return_counts=True) return_counts=True)
return dict(zip(clabels.tolist(), counts.tolist())) return dict(zip(clabels.tolist(), counts.tolist()))
def _initialize_components(self, initializer, num_components): def _initialize_components(self, num_components, initializer):
if isinstance(initializer, ClassAwareInitializer): if isinstance(initializer, ClassAwareInitializer):
_precheck_initializer(initializer) _precheck_initializer(initializer)
_components = initializer.generate(num_components, _components = initializer.generate(num_components,
self.initial_distribution) self.initial_distribution)
self._register_components(_components) self._register_components(_components)
else: else:
super()._initialize_components(initializer, num_components) super()._initialize_components(num_components, initializer)
def add_components(self, initializer, distribution): def add_components(self, distribution, initializer):
_precheck_initializer(initializer) _precheck_initializer(initializer)
# Labels # Labels