Reuse thread pool

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

cvxpylayers/__init__.py

@@ -1 +1 @@
-__version__ = "0.1.2"
+__version__ = "0.1.3"

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,

setup.py

@@ -7,15 +7,15 @@
 
 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",
     packages=find_packages(),
     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",