unit test with multiple omp threads (#14958)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/14958

Test with multiple threads

Reviewed By: jianyuh

Differential Revision: D13394791

fbshipit-source-id: 931a6c3bda15ebc816807e537dd0841c383e7a6f

diff --git a/caffe2/quantization/server/batch_matmul_dnnlowp_op_test.py b/caffe2/quantization/server/batch_matmul_dnnlowp_op_test.py
index 6bfea89..19e0c73 100644 (file)
--- a/caffe2/quantization/server/batch_matmul_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/batch_matmul_dnnlowp_op_test.py
@@ -6,7 +6,7 @@ from itertools import product
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from caffe2.quantization.server import utils as dnnlowp_utils
 from dnnlowp_test_utils import (
     avoid_vpmaddubsw_overflow_fc,
@@ -16,6 +16,7 @@ from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPBatchMatMulOpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/batch_permutation_dnnlowp_op_test.py b/caffe2/quantization/server/batch_permutation_dnnlowp_op_test.py
index 89ecfd0..5990738 100644 (file)
--- a/caffe2/quantization/server/batch_permutation_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/batch_permutation_dnnlowp_op_test.py
@@ -8,6 +8,7 @@ from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPBatchPermutationOpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/channel_shuffle_dnnlowp_op.cc b/caffe2/quantization/server/channel_shuffle_dnnlowp_op.cc
index 6930c73..217cd48 100644 (file)
--- a/caffe2/quantization/server/channel_shuffle_dnnlowp_op.cc
+++ b/caffe2/quantization/server/channel_shuffle_dnnlowp_op.cc
@@ -12,7 +12,7 @@ ChannelShuffleDNNLowPOp<T>::ChannelShuffleDNNLowPOp(
     Workspace* ws)
     : BaseType(operator_def, ws),
       order_(StringToStorageOrder(
-          OperatorBase::GetSingleArgument<std::string>("order", "NCHW"))),
+          this->template GetSingleArgument<std::string>("order", "NCHW"))),
       OP_SINGLE_ARG(int, "group", group_, 1) {
   CAFFE_ENFORCE_NE(order_, StorageOrder::UNKNOWN);
 }
@@ -49,13 +49,11 @@ bool ChannelShuffleDNNLowPOp<T>::RunOnDeviceWithOrderNCHW() {
 #pragma omp parallel for
 #endif
   for (int i = 0; i < N; ++i) {
-    ConstEigenMatrixMap<T> X_mat(X_data, K * HxW, G);
+    ConstEigenMatrixMap<T> X_mat(X_data + i * stride, K * HxW, G);
     for (int j = 0; j < K; ++j) {
-      EigenMatrixMap<T>(Y_data + j * G * HxW, HxW, G) =
+      EigenMatrixMap<T>(Y_data + i * stride + j * G * HxW, HxW, G) =
           X_mat.block(j * HxW, 0, HxW, G);
     }
-    X_data += stride;
-    Y_data += stride;
   }
 
   // Even if there is a pre-chosen quantization parameters for the output,
@@ -95,7 +93,9 @@ bool ChannelShuffleDNNLowPOp<T>::RunOnDeviceWithOrderNHWC() {
     for (auto i = 0; i < X.numel(); i += C) {
       // Transpose each C = GxK matrix
       fbgemm::transpose_4rows(
-          K, (const std::uint8_t*)(X_data + i), (std::uint8_t*)(Y_data + i));
+          K,
+          reinterpret_cast<const std::uint8_t*>(X_data + i),
+          reinterpret_cast<std::uint8_t*>(Y_data + i));
     }
   } else {
 #ifdef _OPENMP

diff --git a/caffe2/quantization/server/channel_shuffle_dnnlowp_op_test.py b/caffe2/quantization/server/channel_shuffle_dnnlowp_op_test.py
index fe11b93..95149aa 100644 (file)
--- a/caffe2/quantization/server/channel_shuffle_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/channel_shuffle_dnnlowp_op_test.py
@@ -8,6 +8,7 @@ from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPChannelShuffleOpsTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/concat_dnnlowp_op_test.py b/caffe2/quantization/server/concat_dnnlowp_op_test.py
index 87787a3..966440c 100644 (file)
--- a/caffe2/quantization/server/concat_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/concat_dnnlowp_op_test.py
@@ -5,12 +5,13 @@ import collections
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from dnnlowp_test_utils import check_quantized_results_close
 from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPConcatOpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/conv_depthwise_dnnlowp_op_test.py b/caffe2/quantization/server/conv_depthwise_dnnlowp_op_test.py
index 4c1953d..70001e1 100644 (file)
--- a/caffe2/quantization/server/conv_depthwise_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/conv_depthwise_dnnlowp_op_test.py
@@ -4,7 +4,7 @@ import collections
 
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from caffe2.quantization.server import utils as dnnlowp_utils
 from dnnlowp_test_utils import (
     check_quantized_results_close,
@@ -17,6 +17,7 @@ from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPOpConvDepthWiseTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/conv_dnnlowp_acc16_op_test.py b/caffe2/quantization/server/conv_dnnlowp_acc16_op_test.py
index cfa82d2..d8d6a79 100644 (file)
--- a/caffe2/quantization/server/conv_dnnlowp_acc16_op_test.py
+++ b/caffe2/quantization/server/conv_dnnlowp_acc16_op_test.py
@@ -5,7 +5,7 @@ import collections
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from caffe2.quantization.server import utils as dnnlowp_utils
 from dnnlowp_test_utils import (
     check_quantized_results_close,
@@ -16,6 +16,7 @@ from hypothesis import assume, given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPOpConvAcc16OpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/conv_dnnlowp_op_test.py b/caffe2/quantization/server/conv_dnnlowp_op_test.py
index 3678780..d197aa7 100644 (file)
--- a/caffe2/quantization/server/conv_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/conv_dnnlowp_op_test.py
@@ -4,7 +4,7 @@ import collections
 
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from caffe2.quantization.server import utils as dnnlowp_utils
 from dnnlowp_test_utils import (
     check_quantized_results_close,
@@ -17,6 +17,7 @@ from hypothesis import assume, given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPOpConvTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/conv_groupwise_dnnlowp_acc16_op_test.py b/caffe2/quantization/server/conv_groupwise_dnnlowp_acc16_op_test.py
index 032c04f..2a227f3 100644 (file)
--- a/caffe2/quantization/server/conv_groupwise_dnnlowp_acc16_op_test.py
+++ b/caffe2/quantization/server/conv_groupwise_dnnlowp_acc16_op_test.py
@@ -5,7 +5,7 @@ import collections
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from caffe2.quantization.server import utils as dnnlowp_utils
 from dnnlowp_test_utils import (
     check_quantized_results_close,
@@ -16,6 +16,7 @@ from hypothesis import assume, given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class GroupWiseDNNLowPOpConvAcc16OpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/conv_groupwise_dnnlowp_op_test.py b/caffe2/quantization/server/conv_groupwise_dnnlowp_op_test.py
index 9ca2461..8e73edd 100644 (file)
--- a/caffe2/quantization/server/conv_groupwise_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/conv_groupwise_dnnlowp_op_test.py
@@ -4,7 +4,7 @@ import collections
 
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from caffe2.python.fb import hardcode_scale_zp
 from caffe2.quantization.server import utils as dnnlowp_utils
 from dnnlowp_test_utils import check_quantized_results_close, generate_conv_inputs
@@ -12,6 +12,7 @@ from hypothesis import assume, given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class GroupWiseDNNLowPOpConvTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/dequantize_dnnlowp_op_test.py b/caffe2/quantization/server/dequantize_dnnlowp_op_test.py
index 8d60dd5..50961de 100644 (file)
--- a/caffe2/quantization/server/dequantize_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/dequantize_dnnlowp_op_test.py
@@ -5,12 +5,13 @@ import collections
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from dnnlowp_test_utils import check_quantized_results_close
 from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPDequantizeOpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/elementwise_add_dnnlowp_op_test.py b/caffe2/quantization/server/elementwise_add_dnnlowp_op_test.py
index 5c24709..d02e66e 100644 (file)
--- a/caffe2/quantization/server/elementwise_add_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/elementwise_add_dnnlowp_op_test.py
@@ -5,12 +5,13 @@ import collections
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from dnnlowp_test_utils import check_quantized_results_close
 from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPAddOpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/elementwise_linear_dnnlowp_op_test.py b/caffe2/quantization/server/elementwise_linear_dnnlowp_op_test.py
index 2d2a135..8b3877e 100644 (file)
--- a/caffe2/quantization/server/elementwise_linear_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/elementwise_linear_dnnlowp_op_test.py
@@ -5,12 +5,13 @@ import collections
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from dnnlowp_test_utils import check_quantized_results_close
 from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPElementwiseLinearOpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/elementwise_mul_dnnlowp_op_test.py b/caffe2/quantization/server/elementwise_mul_dnnlowp_op_test.py
index e122bd1..e63d977 100644 (file)
--- a/caffe2/quantization/server/elementwise_mul_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/elementwise_mul_dnnlowp_op_test.py
@@ -5,12 +5,13 @@ import collections
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from dnnlowp_test_utils import check_quantized_results_close
 from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPMulOpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/elementwise_sum_dnnlowp_op_test.py b/caffe2/quantization/server/elementwise_sum_dnnlowp_op_test.py
index efd1bd0..ec752b0 100644 (file)
--- a/caffe2/quantization/server/elementwise_sum_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/elementwise_sum_dnnlowp_op_test.py
@@ -5,12 +5,13 @@ import collections
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from dnnlowp_test_utils import check_quantized_results_close
 from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPOpSumOpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/fully_connected_dnnlowp_acc16_op_test.py b/caffe2/quantization/server/fully_connected_dnnlowp_acc16_op_test.py
index ce3598e..ba1a8c8 100644 (file)
--- a/caffe2/quantization/server/fully_connected_dnnlowp_acc16_op_test.py
+++ b/caffe2/quantization/server/fully_connected_dnnlowp_acc16_op_test.py
@@ -5,13 +5,14 @@ import collections
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from caffe2.quantization.server import utils as dnnlowp_utils
 from dnnlowp_test_utils import check_quantized_results_close
 from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPFullyConnectedAcc16OpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/fully_connected_dnnlowp_op_test.py b/caffe2/quantization/server/fully_connected_dnnlowp_op_test.py
index edbf38e..e08b635 100644 (file)
--- a/caffe2/quantization/server/fully_connected_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/fully_connected_dnnlowp_op_test.py
@@ -5,7 +5,7 @@ import collections
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from caffe2.quantization.server import utils as dnnlowp_utils
 from dnnlowp_test_utils import (
     avoid_vpmaddubsw_overflow_fc,
@@ -15,6 +15,7 @@ from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPFullyConnectedOpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/fully_connected_fp16_test.py b/caffe2/quantization/server/fully_connected_fp16_test.py
index 4bcd2f0..ad130f3 100644 (file)
--- a/caffe2/quantization/server/fully_connected_fp16_test.py
+++ b/caffe2/quantization/server/fully_connected_fp16_test.py
@@ -5,11 +5,12 @@ import collections
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 def mse(x, xh):

diff --git a/caffe2/quantization/server/fully_connected_rowwise_dnnlowp_op_test.py b/caffe2/quantization/server/fully_connected_rowwise_dnnlowp_op_test.py
index 550ed7b..65b17fc 100644 (file)
--- a/caffe2/quantization/server/fully_connected_rowwise_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/fully_connected_rowwise_dnnlowp_op_test.py
@@ -5,7 +5,7 @@ import collections
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from caffe2.quantization.server import utils as dnnlowp_utils
 from dnnlowp_test_utils import (
     avoid_vpmaddubsw_overflow_fc,
@@ -15,6 +15,7 @@ from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class RowWiseDNNLowPFullyConnectedOpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/gather_dnnlowp_op_test.py b/caffe2/quantization/server/gather_dnnlowp_op_test.py
index 8abc3f5..0c1d7e9 100644 (file)
--- a/caffe2/quantization/server/gather_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/gather_dnnlowp_op_test.py
@@ -5,12 +5,13 @@ import collections
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from dnnlowp_test_utils import check_quantized_results_close
 from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPGatherOpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/group_norm_dnnlowp_op_test.py b/caffe2/quantization/server/group_norm_dnnlowp_op_test.py
index d386ea1..bd0bfff 100644 (file)
--- a/caffe2/quantization/server/group_norm_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/group_norm_dnnlowp_op_test.py
@@ -5,13 +5,14 @@ import collections
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from caffe2.quantization.server import utils as dnnlowp_utils
 from dnnlowp_test_utils import check_quantized_results_close
 from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPOpGroupNormTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/lstm_unit_dnnlowp_op_test.py b/caffe2/quantization/server/lstm_unit_dnnlowp_op_test.py
index 0a1e78a..aaa9086 100644 (file)
--- a/caffe2/quantization/server/lstm_unit_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/lstm_unit_dnnlowp_op_test.py
@@ -5,11 +5,12 @@ import collections
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPLSTMUnitOpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/pool_dnnlowp_op_test.py b/caffe2/quantization/server/pool_dnnlowp_op_test.py
index 49fde98..6546afb 100644 (file)
--- a/caffe2/quantization/server/pool_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/pool_dnnlowp_op_test.py
@@ -5,12 +5,13 @@ import collections
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from dnnlowp_test_utils import check_quantized_results_close
 from hypothesis import assume, given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPOpPoolTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/quantize_dnnlowp_op_test.py b/caffe2/quantization/server/quantize_dnnlowp_op_test.py
index 85efcae..97a656c 100644 (file)
--- a/caffe2/quantization/server/quantize_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/quantize_dnnlowp_op_test.py
@@ -3,11 +3,12 @@ from __future__ import absolute_import, division, print_function, unicode_litera
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPQuantizeOpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/relu_dnnlowp_op_test.py b/caffe2/quantization/server/relu_dnnlowp_op_test.py
index a360e47..0f893dd 100644 (file)
--- a/caffe2/quantization/server/relu_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/relu_dnnlowp_op_test.py
@@ -5,13 +5,14 @@ from __future__ import unicode_literals
 
 import numpy as np
 import caffe2.python.hypothesis_test_util as hu
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from hypothesis import given
 import hypothesis.strategies as st
 import collections
 from dnnlowp_test_utils import check_quantized_results_close
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPReluOpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/resize_nearest_dnnlowp_op_test.py b/caffe2/quantization/server/resize_nearest_dnnlowp_op_test.py
index f11f2a1..519a23e 100644 (file)
--- a/caffe2/quantization/server/resize_nearest_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/resize_nearest_dnnlowp_op_test.py
@@ -8,6 +8,7 @@ from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPResizeNearestOpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/sigmoid_dnnlowp_op_test.py b/caffe2/quantization/server/sigmoid_dnnlowp_op_test.py
index dd73a8d..aa97b32 100644 (file)
--- a/caffe2/quantization/server/sigmoid_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/sigmoid_dnnlowp_op_test.py
@@ -5,11 +5,12 @@ import collections
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPSigmoidOpTest(hu.HypothesisTestCase):

diff --git a/caffe2/quantization/server/tanh_dnnlowp_op_test.py b/caffe2/quantization/server/tanh_dnnlowp_op_test.py
index 6037f63..636e40a 100644 (file)
--- a/caffe2/quantization/server/tanh_dnnlowp_op_test.py
+++ b/caffe2/quantization/server/tanh_dnnlowp_op_test.py
@@ -5,11 +5,12 @@ import collections
 import caffe2.python.hypothesis_test_util as hu
 import hypothesis.strategies as st
 import numpy as np
-from caffe2.python import core, dyndep
+from caffe2.python import core, dyndep, workspace
 from hypothesis import given
 
 
 dyndep.InitOpsLibrary("//caffe2/caffe2/quantization/server:dnnlowp_ops")
+workspace.GlobalInit(["caffe2", "--caffe2_omp_num_threads=11"])
 
 
 class DNNLowPTanhOpTest(hu.HypothesisTestCase):