diff --git a/README.md b/README.md index ebddb42..ec385bd 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,7 @@ cvxpylayers has the following dependencies: * [NumPy](https://pypi.org/project/numpy/) * [CVXPY](https://github.com/cvxgrp/cvxpy) >= 1.1.a1 * [TensorFlow](https://tensorflow.org) >= 2.0 or [PyTorch](https://pytorch.org) >= 1.0 -* [diffcp](https://github.com/cvxgrp/diffcp) >= 1.0.13 +* [diffcp](https://github.com/cvxgrp/diffcp) >= 1.0.14 ## Usage Below are usage examples of our PyTorch and TensorFlow layers. Note that diff --git a/cvxpylayers/__init__.py b/cvxpylayers/__init__.py index b3f4756..ae73625 100644 --- a/cvxpylayers/__init__.py +++ b/cvxpylayers/__init__.py @@ -1 +1 @@ -__version__ = "0.1.2" +__version__ = "0.1.3" diff --git a/cvxpylayers/torch/cvxpylayer.py b/cvxpylayers/torch/cvxpylayer.py index e84d777..be0447a 100644 --- a/cvxpylayers/torch/cvxpylayer.py +++ b/cvxpylayers/torch/cvxpylayer.py @@ -4,6 +4,8 @@ from cvxpy.reductions.solvers.conic_solvers.scs_conif import \ dims_to_solver_dict import numpy as np +import multiprocessing as mp +from multiprocessing.pool import ThreadPool try: import torch @@ -109,6 +111,34 @@ def forward(self, *params, solver_args={}): 'parameter; received %d tensors, expected %d' % ( len(params), len(self.param_ids))) info = {} + + n_jobs_forward = solver_args.get('n_jobs_forward', -1) + n_jobs_backward = solver_args.get('n_jobs_backward', -1) + if n_jobs_forward == -1: + n_jobs_forward = mp.cpu_count() + if n_jobs_backward == -1: + n_jobs_backward = mp.cpu_count() + + if n_jobs_forward != 1: + if hasattr(self, "pool_forward"): + forward_threads = len(self.pool_forward._pool) + if forward_threads != n_jobs_forward: + self.pool_forward.close() + self.pool_forward = ThreadPool(n_jobs_forward) + else: + self.pool_forward = ThreadPool(n_jobs_forward) + solver_args["pool_forward"] = self.pool_forward + + if n_jobs_backward != 1: + if hasattr(self, "pool_backward"): + backward_threads = len(self.pool_backward._pool) + if backward_threads != n_jobs_backward: + self.pool_backward.close() + self.pool_backward = ThreadPool(n_jobs_backward) + else: + self.pool_backward = ThreadPool(n_jobs_backward) + solver_args["pool_backward"] = self.pool_backward + f = _CvxpyLayerFn( param_order=self.param_order, param_ids=self.param_ids, diff --git a/setup.py b/setup.py index c469854..7be5e85 100644 --- a/setup.py +++ b/setup.py @@ -7,7 +7,7 @@ setup( name="cvxpylayers", - version="0.1.2", + version="0.1.3", description="Differentiable convex optimization layers", long_description=long_description, long_description_content_type="text/markdown", @@ -15,7 +15,7 @@ install_requires=[ "numpy >= 1.15", "scipy >= 1.1.0", - "diffcp >= 1.0.13", + "diffcp >= 1.0.14", "cvxpy >= 1.1.0a1"], license="Apache License, Version 2.0", url="https://github.com/cvxgrp/cvxpylayers",