diff --git a/learnware/specification/rkme.py b/learnware/specification/rkme.py index 2ca749f..33440b9 100644 --- a/learnware/specification/rkme.py +++ b/learnware/specification/rkme.py @@ -4,20 +4,117 @@ import torch import faiss import json import codecs +import random import numpy as np -from .base import BaseStatSpecification -from .utils import setup_seed, choose_device, torch_rbf_kernel, solve_qp +from cvxopt import solvers, matrix from typing import Tuple, Any, List, Union, Dict -from learnware.config import C + +from .base import BaseStatSpecification class RKMESpecification: pass +def setup_seed(seed): + """ + Fix a random seed for addressing reproducibility issues. + + Parameters + ---------- + seed : int + Random seed for torch, torch.cuda, numpy, random and cudnn libraries. + """ + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + np.random.seed(seed) + random.seed(seed) + torch.backends.cudnn.deterministic = True + + +def choose_device(cuda_idx=-1): + """ + Let users choose compuational device between CPU or GPU. + + Parameters + ---------- + cuda_idx : int, optional + GPU index, by default -1 which stands for using CPU instead. + + Returns + ------- + torch.device + A torch.device object + """ + if cuda_idx != -1: + device = torch.device(f"cuda:{cuda_idx}") + else: + device = torch.device("cpu") + return device + + +def torch_rbf_kernel(x1, x2, gamma) -> torch.Tensor: + """ + Use pytorch to compute rbf_kernel function at faster speed. + + Parameters + ---------- + x1 : torch.Tensor + First vector in the rbf_kernel + x2 : torch.Tensor + Second vector in the rbf_kernel + gamma : float + Bandwidth in gaussian kernel + + Returns + ------- + torch.Tensor + The computed rbf_kernel value at x1, x2. + """ + x1 = x1.double() + x2 = x2.double() + X12norm = torch.sum(x1 ** 2, 1, keepdim=True) - 2 * x1 @ x2.T + torch.sum(x2 ** 2, 1, keepdim=True).T + return torch.exp(-X12norm * gamma) + + +def solve_qp(K: np.ndarray, C: np.ndarray): + """ + 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 + + Parameters + ---------- + K : np.ndarray + Parameter in the quadratic term. + C : np.ndarray + Parameter in the linear term. + + Returns + ------- + torch.tensor + Solution to the QP problem. + """ + n = K.shape[0] + P = matrix(K.cpu().numpy()) + q = matrix(-C.cpu().numpy()) + G = matrix(-np.eye(n)) + h = matrix(np.zeros((n, 1))) + A = matrix(np.ones((1, n))) + b = matrix(np.ones((1, 1))) + + solvers.options["show_progress"] = False + sol = solvers.qp(P, q, G, h, A, b) # Requires the sum of x to be 1 + # sol = solvers.qp(P, q, G, h) # Otherwise + w = np.array(sol["x"]) + w = torch.from_numpy(w).reshape(-1) + + return w + + class RKMESpecification(BaseStatSpecification): - """Reduced-set Kernel Mean Embedding(RKME) Specification - + """Reduced-set Kernel Mean Embedding (RKME) Specification """ def __init__(self, gamma: float = 0.1, cuda_idx: int = -1): @@ -30,8 +127,8 @@ class RKMESpecification(BaseStatSpecification): cuda_idx : int A flag indicating whether use CUDA during RKME computation. -1 indicates CUDA not used. """ - self.Z = [] - self.beta = [] + self.z = None + self.beta = None self.gamma = gamma self.num_points = 0 self.cuda_idx = cuda_idx @@ -39,28 +136,34 @@ class RKMESpecification(BaseStatSpecification): self.device = choose_device(cuda_idx=cuda_idx) setup_seed(0) - def get_beta(self) -> torch.tensor: + def get_beta(self) -> np.ndarray: """Move beta(RKME weights) back to memory accessible to the CPU. Returns ------- - torch.tensor + np.ndarray A copy of beta in CPU memory. """ - return self.beta.detach().cpu() + return self.beta.detach().cpu().numpy() - def get_z(self) -> torch.tensor: + def get_z(self) -> np.ndarray: """Move z(RKME reduced set points) back to memory accessible to the CPU. Returns ------- - torch.tensor + np.ndarray A copy of z in CPU memory. """ - return self.z.detach().cpu() + return self.z.detach().cpu().numpy() def generate_stat_spec_from_data( - self, X: np.ndarray, K: int, step_size: float, steps: int, reduce: bool = True, nonnegative_beta: bool = False + self, + X: np.ndarray, + K: int = 100, + step_size: float = 0.1, + steps: int = 3, + nonnegative_beta: bool = True, + reduce: bool = True, ): """Construct reduced set from raw dataset using iterative optimization. @@ -74,14 +177,21 @@ class RKMESpecification(BaseStatSpecification): Step size for gradient descent in the iterative optimization. steps : int Total rounds in the iterative optimization. - reduce : bool, optional - Whether shrink original data to a smaller set, by default True nonnegative_beta : bool, optional True if weights for the reduced set are intended to be kept non-negative, by default False. + reduce : bool, optional + Whether shrink original data to a smaller set, by default True """ alpha = None self.num_points = X.shape[0] - + + # fill np.nan + X_nan = np.isnan(X) + if X_nan.max() == 1: + for col in range(X.shape[1]): + col_mean = np.nanmean(X[:, col]) + X[:, col] = np.where(X_nan[:, col], col_mean, X[:, col]) + if not reduce: self.z = X self.beta = 1 / self.num_points * np.ones(self.num_points) @@ -98,7 +208,7 @@ class RKMESpecification(BaseStatSpecification): self._update_z(alpha, X, step_size) self._update_beta(X, nonnegative_beta) - def _init_z_by_faiss(self, X: Any, K: int): + def _init_z_by_faiss(self, X: Union[np.ndarray, torch.tensor], K: int): """Intialize Z by faiss clustering. Parameters @@ -108,6 +218,7 @@ class RKMESpecification(BaseStatSpecification): K : int Size of the construced reduced set. """ + X = X.astype("float32") numDim = X.shape[1] kmeans = faiss.Kmeans(numDim, K, niter=100, verbose=False) kmeans.train(X) @@ -185,7 +296,7 @@ class RKMESpecification(BaseStatSpecification): Z = Z - step_size * grad_Z self.z = Z - def eval_Phi(self, Phi2: RKMESpecification) -> float: + def inner_prod(self, Phi2: RKMESpecification) -> float: """Compute the inner product between two RKME specifications Parameters @@ -206,7 +317,7 @@ class RKMESpecification(BaseStatSpecification): v = torch.sum(torch_rbf_kernel(Z1, Z2, self.gamma) * (beta_1.T @ beta_2)) return float(v) - def MMD(self, Phi2: RKMESpecification, omit_term1: bool = False) -> float: + def dist(self, Phi2: RKMESpecification, omit_term1: bool = False) -> float: """Compute the Maximum-Mean-Discrepancy(MMD) between two RKME specifications Parameters @@ -219,15 +330,12 @@ class RKMESpecification(BaseStatSpecification): if omit_term1: term1 = 0 else: - term1 = self.eval_Phi(self) - term2 = self.eval_Phi(Phi2) - term3 = Phi2.eval_Phi(Phi2) + term1 = self.inner_prod(self) + term2 = self.inner_prod(Phi2) + term3 = Phi2.inner_prod(Phi2) return float(term1 - 2 * term2 + term3) - def generate_stat_spec_from_data(self, X: np.ndarray): - return super().generate_stat_spec_from_data(X) - def save(self, filepath: str): """Save the computed RKME specification to a specified path in JSON format. @@ -236,7 +344,7 @@ class RKMESpecification(BaseStatSpecification): filepath : str The specified saving path. """ - save_path = os.path.join(C.specification_path, f"{filepath}.json") + save_path = filepath rkme_to_save = copy.deepcopy(self.__dict__) if torch.is_tensor(rkme_to_save["z"]): rkme_to_save["z"] = rkme_to_save["z"].detach().cpu().numpy() @@ -263,7 +371,7 @@ class RKMESpecification(BaseStatSpecification): True if the RKME is loaded successfully. """ # Load JSON file: - load_path = os.path.join(C.specification_path, f"{filepath}.json") + load_path = filepath if os.path.exists(load_path): obj_text = codecs.open(load_path, "r", encoding="utf-8").read() rkme_load = json.loads(obj_text) diff --git a/learnware/specification/utils.py b/learnware/specification/utils.py index cabeb6f..e678d5c 100644 --- a/learnware/specification/utils.py +++ b/learnware/specification/utils.py @@ -1,105 +1,52 @@ import numpy as np -import torch -import random -from cvxopt import solvers, matrix -from .base import BaseStatSpecification - - -def setup_seed(seed): - """ - Fix a random seed for addressing reproducibility issues. - - Parameters - ---------- - seed : int - Random seed for torch, torch.cuda, numpy, random and cudnn libraries. - """ - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - np.random.seed(seed) - random.seed(seed) - torch.backends.cudnn.deterministic = True - -def choose_device(cuda_idx=-1): +from .base import BaseStatSpecification +from .rkme import RKMESpecification + + +def generate_rkme_spec( + X: np.ndarray, + gamma: float = 0.1, + K: int = 100, + step_size: float = 0.1, + steps: int = 3, + nonnegative_beta: bool = True, + reduce: bool = True, + cuda_idx: int = -1, +) -> RKMESpecification: """ - Let users choose compuational device between CPU or GPU. - - Parameters - ---------- - cuda_idx : int, optional - GPU index, by default -1 which stands for using CPU instead. - - Returns - ------- - torch.device - A torch.device object - """ - if cuda_idx != -1: - device = torch.device(f"cuda:{cuda_idx}") - else: - device = torch.device("cpu") - return device + Interface for users to generate Reduced-set Kernel Mean Embedding (RKME) specification. + Return a RKMESpecification object, use .save() method to save as json file. -def torch_rbf_kernel(x1, x2, gamma) -> torch.Tensor: - """ - Use pytorch to compute rbf_kernel function at faster speed. - Parameters ---------- - x1 : torch.Tensor - First vector in the rbf_kernel - x2 : torch.Tensor - Second vector in the rbf_kernel + X : np.ndarray + Raw data in np.ndarray format. + Size of array: (n*d) gamma : float - Bandwidth in gaussian kernel - - Returns - ------- - torch.Tensor - The computed rbf_kernel value at x1, x2. - """ - x1 = x1.double() - x2 = x2.double() - X12norm = torch.sum(x1 ** 2, 1, keepdim=True) - 2 * x1 @ x2.T + torch.sum(x2 ** 2, 1, keepdim=True).T - return torch.exp(-X12norm * gamma) - - -def solve_qp(K: np.ndarray, C: np.ndarray): - """ - 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 - - Parameters - ---------- - K : np.ndarray - Parameter in the quadratic term. - C : np.ndarray - Parameter in the linear term. + Bandwidth in gaussian kernel, by default 0.1. + K : int + Size of the construced reduced set. + step_size : float + Step size for gradient descent in the iterative optimization. + steps : int + Total rounds in the iterative optimization. + nonnegative_beta : bool, optional + True if weights for the reduced set are intended to be kept non-negative, by default False. + reduce : bool, optional + Whether shrink original data to a smaller set, by default True + cuda_idx : int + A flag indicating whether use CUDA during RKME computation. -1 indicates CUDA not used. Returns ------- - torch.tensor - Solution to the QP problem. + RKMESpecification + A RKMESpecification object """ - n = K.shape[0] - P = matrix(K.cpu().numpy()) - q = matrix(-C.cpu().numpy()) - G = matrix(-np.eye(n)) - h = matrix(np.zeros((n, 1))) - A = matrix(np.ones((1, n))) - b = matrix(np.ones((1, 1))) - - solvers.options["show_progress"] = False - sol = solvers.qp(P, q, G, h, A, b) # Requires the sum of x to be 1 - # sol = solvers.qp(P, q, G, h) # Otherwise - w = np.array(sol["x"]) - w = torch.from_numpy(w).reshape(-1) - - return w + rkme_spec = RKMESpecification(gamma=gamma, cuda_idx=cuda_idx) + rkme_spec.generate_stat_spec_from_data(X, K, step_size, steps, nonnegative_beta, reduce) + return rkme_spec def generate_stat_spec(X: np.ndarray) -> BaseStatSpecification: