diff --git a/miplearn/components/component.py b/miplearn/components/component.py index f445a63..986fb3a 100644 --- a/miplearn/components/component.py +++ b/miplearn/components/component.py @@ -11,6 +11,8 @@ from miplearn.features import Sample from miplearn.instance.base import Instance from miplearn.types import LearningSolveStats +from p_tqdm import p_umap + if TYPE_CHECKING: from miplearn.solvers.learning import LearningSolver @@ -159,7 +161,6 @@ class Component(EnforceOverrides): self, instance: Optional[Instance], sample: Sample, - pre: Optional[List[Any]] = None, ) -> Tuple[Dict, Dict]: """ Returns a pair of x and y dictionaries containing, respectively, the matrices @@ -168,6 +169,9 @@ class Component(EnforceOverrides): """ pass + def pre_fit(self, pre: List[Any]): + pass + def user_cut_cb( self, solver: "LearningSolver", @@ -183,6 +187,7 @@ class Component(EnforceOverrides): def fit_multiple( components: Dict[str, "Component"], instances: List[Instance], + n_jobs: int = 1, ) -> None: def _pre_sample_xy(instance: Instance) -> Dict: pre_instance: Dict = {} @@ -195,7 +200,17 @@ class Component(EnforceOverrides): instance.free() return pre_instance - def _sample_xy(instance: Instance, pre: Dict) -> Tuple[Dict, Dict]: + pre = p_umap(_pre_sample_xy, instances, num_cpus=n_jobs) + pre_combined: Dict = {} + for (cname, comp) in components.items(): + pre_combined[cname] = [] + for p in pre: + pre_combined[cname].extend(p[cname]) + + for (cname, comp) in components.items(): + comp.pre_fit(pre_combined[cname]) + + def _sample_xy(instance: Instance) -> Tuple[Dict, Dict]: x_instance: Dict = {} y_instance: Dict = {} for (cname, comp) in components.items(): @@ -206,7 +221,7 @@ class Component(EnforceOverrides): for (cname, comp) in components.items(): x = x_instance[cname] y = y_instance[cname] - x_sample, y_sample = comp.sample_xy(instance, sample, pre[cname]) + x_sample, y_sample = comp.sample_xy(instance, sample) for cat in x_sample.keys(): if cat not in x: x[cat] = [] @@ -216,15 +231,7 @@ class Component(EnforceOverrides): instance.free() return x_instance, y_instance - pre = [_pre_sample_xy(instance) for instance in instances] - - pre_combined: Dict = {} - for (cname, comp) in components.items(): - pre_combined[cname] = [] - for p in pre: - pre_combined[cname].extend(p[cname]) - - xy = [_sample_xy(instances, pre_combined) for instances in instances] + xy = p_umap(_sample_xy, instances) for (cname, comp) in components.items(): x_comp: Dict = {} diff --git a/miplearn/components/dynamic_common.py b/miplearn/components/dynamic_common.py index 561becf..5b63a9b 100644 --- a/miplearn/components/dynamic_common.py +++ b/miplearn/components/dynamic_common.py @@ -89,16 +89,18 @@ class DynamicConstraintsComponent(Component): self, instance: Optional[Instance], sample: Sample, - pre: Optional[List[Any]] = None, ) -> Tuple[Dict, Dict]: + x, y, _ = self.sample_xy_with_cids(instance, sample) + return x, y + + @overrides + def pre_fit(self, pre: List[Any]) -> None: assert pre is not None known_cids: Set = set() for cids in pre: known_cids |= cids self.known_cids.clear() self.known_cids.extend(sorted(known_cids)) - x, y, _ = self.sample_xy_with_cids(instance, sample) - return x, y def sample_predict( self, diff --git a/miplearn/components/dynamic_lazy.py b/miplearn/components/dynamic_lazy.py index ab1d412..d0843d6 100644 --- a/miplearn/components/dynamic_lazy.py +++ b/miplearn/components/dynamic_lazy.py @@ -108,9 +108,12 @@ class DynamicLazyConstraintsComponent(Component): self, instance: Optional[Instance], sample: Sample, - pre: Optional[List[Any]] = None, ) -> Tuple[Dict, Dict]: - return self.dynamic.sample_xy(instance, sample, pre=pre) + return self.dynamic.sample_xy(instance, sample) + + @overrides + def pre_fit(self, pre: List[Any]) -> None: + self.dynamic.pre_fit(pre) def sample_predict( self, diff --git a/miplearn/components/dynamic_user_cuts.py b/miplearn/components/dynamic_user_cuts.py index ebf7214..ec12dc0 100644 --- a/miplearn/components/dynamic_user_cuts.py +++ b/miplearn/components/dynamic_user_cuts.py @@ -101,9 +101,12 @@ class UserCutsComponent(Component): self, instance: "Instance", sample: Sample, - pre: Optional[List[Any]] = None, ) -> Tuple[Dict, Dict]: - return self.dynamic.sample_xy(instance, sample, pre=pre) + return self.dynamic.sample_xy(instance, sample) + + @overrides + def pre_fit(self, pre: List[Any]) -> None: + self.dynamic.pre_fit(pre) def sample_predict( self, diff --git a/miplearn/components/objective.py b/miplearn/components/objective.py index a226bc7..2fc5afd 100644 --- a/miplearn/components/objective.py +++ b/miplearn/components/objective.py @@ -76,7 +76,6 @@ class ObjectiveValueComponent(Component): self, _: Optional[Instance], sample: Sample, - pre: Optional[List[Any]] = None, ) -> Tuple[Dict[Hashable, List[List[float]]], Dict[Hashable, List[List[float]]]]: # Instance features assert sample.after_load is not None diff --git a/miplearn/components/primal.py b/miplearn/components/primal.py index 15359e3..00c676e 100644 --- a/miplearn/components/primal.py +++ b/miplearn/components/primal.py @@ -145,7 +145,6 @@ class PrimalSolutionComponent(Component): self, _: Optional[Instance], sample: Sample, - pre: Optional[List[Any]] = None, ) -> Tuple[Dict[Category, List[List[float]]], Dict[Category, List[List[float]]]]: x: Dict = {} y: Dict = {} diff --git a/miplearn/components/static_lazy.py b/miplearn/components/static_lazy.py index 6858e5f..ebf204b 100644 --- a/miplearn/components/static_lazy.py +++ b/miplearn/components/static_lazy.py @@ -154,7 +154,6 @@ class StaticLazyConstraintsComponent(Component): self, _: Optional[Instance], sample: Sample, - pre: Optional[List[Any]] = None, ) -> Tuple[Dict[Hashable, List[List[float]]], Dict[Hashable, List[List[float]]]]: x, y, __ = self._sample_xy_with_cids(sample) return x, y diff --git a/miplearn/solvers/learning.py b/miplearn/solvers/learning.py index 8f18ba2..9d41b93 100644 --- a/miplearn/solvers/learning.py +++ b/miplearn/solvers/learning.py @@ -391,11 +391,19 @@ class LearningSolver: self._restore_miplearn_logger() return stats - def fit(self, training_instances: List[Instance]) -> None: + def fit( + self, + training_instances: List[Instance], + n_jobs: int = 1, + ) -> None: if len(training_instances) == 0: logger.warning("Empty list of training instances provided. Skipping.") return - Component.fit_multiple(self.components, training_instances) + Component.fit_multiple( + self.components, + training_instances, + n_jobs=n_jobs, + ) def _add_component(self, component: Component) -> None: name = component.__class__.__name__ diff --git a/tests/components/test_dynamic_lazy.py b/tests/components/test_dynamic_lazy.py index b96c03f..42ae03f 100644 --- a/tests/components/test_dynamic_lazy.py +++ b/tests/components/test_dynamic_lazy.py @@ -87,6 +87,7 @@ def training_instances() -> List[Instance]: def test_sample_xy(training_instances: List[Instance]) -> None: comp = DynamicLazyConstraintsComponent() + comp.pre_fit([{"c1", "c2", "c3", "c4"}]) x_expected = { "type-a": [[5.0, 1.0, 2.0, 3.0], [5.0, 4.0, 5.0, 6.0]], "type-b": [[5.0, 1.0, 2.0], [5.0, 3.0, 4.0]], @@ -98,7 +99,6 @@ def test_sample_xy(training_instances: List[Instance]) -> None: x_actual, y_actual = comp.sample_xy( training_instances[0], training_instances[0].samples[0], - pre=[{"c1", "c2", "c3", "c4"}], ) assert_equals(x_actual, x_expected) assert_equals(y_actual, y_expected)