diff --git a/learnware/specification/rkme.py b/learnware/specification/rkme.py index 33440b9..1f5869e 100644 --- a/learnware/specification/rkme.py +++ b/learnware/specification/rkme.py @@ -1,4 +1,5 @@ import os +import mkl import copy import torch import faiss @@ -11,14 +12,15 @@ from typing import Tuple, Any, List, Union, Dict from .base import BaseStatSpecification +mkl.get_max_threads() + class RKMESpecification: pass def setup_seed(seed): - """ - Fix a random seed for addressing reproducibility issues. + """Fix a random seed for addressing reproducibility issues. Parameters ---------- @@ -33,8 +35,7 @@ def setup_seed(seed): def choose_device(cuda_idx=-1): - """ - Let users choose compuational device between CPU or GPU. + """Let users choose compuational device between CPU or GPU. Parameters ---------- @@ -54,8 +55,7 @@ def choose_device(cuda_idx=-1): def torch_rbf_kernel(x1, x2, gamma) -> torch.Tensor: - """ - Use pytorch to compute rbf_kernel function at faster speed. + """Use pytorch to compute rbf_kernel function at faster speed. Parameters ---------- @@ -78,8 +78,7 @@ def torch_rbf_kernel(x1, x2, gamma) -> torch.Tensor: def solve_qp(K: np.ndarray, C: np.ndarray): - """ - Solver for the following quadratic programming(QP) problem: + """Solver for the following quadratic programming(QP) problem: - min 1/2 x^T K x - C^T x s.t 1^T x - 1 = 0 - I x <= 0