diff --git a/fn_gen/ones_t/0/distortion.png b/fn_gen/ones_t/0/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..d606918fdafc9bfa78e269d336f6434353ba9808 Binary files /dev/null and b/fn_gen/ones_t/0/distortion.png differ diff --git a/fn_gen/ones_t/0/expressions.txt b/fn_gen/ones_t/0/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..3758ee2a62aa8d95c3b7da1dd3fafa11b027ad9b --- /dev/null +++ b/fn_gen/ones_t/0/expressions.txt @@ -0,0 +1,2 @@ +cosh(_0*x)/_s +log(_s*x - sqrt(_s**2*x**2 - 1))/_0 \ No newline at end of file diff --git a/fn_gen/ones_t/0/fn.py b/fn_gen/ones_t/0/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..40d6bf782737d3d25cd9e1e6c01cdcc75d7cb47d --- /dev/null +++ b/fn_gen/ones_t/0/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.cosh((params['_0'] * x))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.log(domain_guard(((torch.tensor(-1) * torch.sqrt(domain_guard((torch.tensor(-1) + (guarded_torch_power(params['_s'], torch.tensor(2)) * guarded_torch_power(x, torch.tensor(2)))), min=0.1, nan=0.1))) + (params['_s'] * x)), min=1e-5, nan=1e-5))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_ones(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.cosh((_0 * x))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.log(np_domain_guard(((np.array(-1) * np.sqrt(np_domain_guard((np.array(-1) + (np_guarded_power(_s, np.array(2)) * np_guarded_power(x, np.array(2)))), min=0.1, nan=0.1))) + (_s * x)), min=1e-5, nan=1e-5))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/0/loss.png b/fn_gen/ones_t/0/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..0c2b69f9767b07b0d6108be59cffdd2186c0f168 Binary files /dev/null and b/fn_gen/ones_t/0/loss.png differ diff --git a/fn_gen/ones_t/0/quantization.png b/fn_gen/ones_t/0/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..6f9b2aac4f6198b27e2d5a14bc151d4997fbb3c8 Binary files /dev/null and b/fn_gen/ones_t/0/quantization.png differ diff --git a/fn_gen/ones_t/1/distortion.png b/fn_gen/ones_t/1/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..3904b84a3f75d52b27c01396a77f221cb5129372 Binary files /dev/null and b/fn_gen/ones_t/1/distortion.png differ diff --git a/fn_gen/ones_t/1/expressions.txt b/fn_gen/ones_t/1/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..c7b68c388fdf6e1b6e2be8076f1d4b8d7bcef4f9 --- /dev/null +++ b/fn_gen/ones_t/1/expressions.txt @@ -0,0 +1,2 @@ +(_0*x)**(1/3)/_s +_s**3*x**3/_0 \ No newline at end of file diff --git a/fn_gen/ones_t/1/fn.py b/fn_gen/ones_t/1/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..7b887903efb53514b16abc36e38986ae62fba0ef --- /dev/null +++ b/fn_gen/ones_t/1/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * guarded_torch_power((params['_0'] * x), 1 / 3)) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * guarded_torch_power(params['_s'], torch.tensor(3)) * guarded_torch_power(x, torch.tensor(3))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_ones(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np_guarded_power((_0 * x), 1 / 3)) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np_guarded_power(_s, np.array(3)) * np_guarded_power(x, np.array(3))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/1/loss.png b/fn_gen/ones_t/1/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..29c03e83acbb1cc181a3cbc0730329b3bd93ec2a Binary files /dev/null and b/fn_gen/ones_t/1/loss.png differ diff --git a/fn_gen/ones_t/1/quantization.png b/fn_gen/ones_t/1/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..ff5ef16dd9667f4f1f192ec90a6574bbaf01557e Binary files /dev/null and b/fn_gen/ones_t/1/quantization.png differ diff --git a/fn_gen/ones_t/10/distortion.png b/fn_gen/ones_t/10/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..9829d8f9e8a27990f7d413d2e6970b823dcf952f Binary files /dev/null and b/fn_gen/ones_t/10/distortion.png differ diff --git a/fn_gen/ones_t/10/expressions.txt b/fn_gen/ones_t/10/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..576ec6a351e26f9982eb17e394804ca906d4b067 --- /dev/null +++ b/fn_gen/ones_t/10/expressions.txt @@ -0,0 +1,2 @@ +acos(_0*x)/_s +cos(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/ones_t/10/fn.py b/fn_gen/ones_t/10/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..989eb2944c80b4420ff4650eb090c26aa89347bb --- /dev/null +++ b/fn_gen/ones_t/10/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.acos(domain_guard((params['_0'] * x), min=-0.99999, max=0.99999, nan=0))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.cos((params['_s'] * x))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_ones(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.arccos(np_domain_guard((_0 * x), min=-0.99999, max=0.99999, nan=0))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.cos((_s * x))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/10/loss.png b/fn_gen/ones_t/10/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..260d944bf798947cbb4b27032e19a0f20a2ad72b Binary files /dev/null and b/fn_gen/ones_t/10/loss.png differ diff --git a/fn_gen/ones_t/10/quantization.png b/fn_gen/ones_t/10/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..b1c5d9c1d0cc5d693ee8680f52efb5e07b966f27 Binary files /dev/null and b/fn_gen/ones_t/10/quantization.png differ diff --git a/fn_gen/ones_t/11/distortion.png b/fn_gen/ones_t/11/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..33fe6a5916084fc8773c298c7d5bb0775d83c146 Binary files /dev/null and b/fn_gen/ones_t/11/distortion.png differ diff --git a/fn_gen/ones_t/11/expressions.txt b/fn_gen/ones_t/11/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..e8458af52eb4cfce21cf8459f3c454003cd78158 --- /dev/null +++ b/fn_gen/ones_t/11/expressions.txt @@ -0,0 +1,2 @@ +sqrt(_0*x)/_s +_s**2*x**2/_0 \ No newline at end of file diff --git a/fn_gen/ones_t/11/fn.py b/fn_gen/ones_t/11/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..1194765af3d61c849ca13a3530ad4a8c634140ab --- /dev/null +++ b/fn_gen/ones_t/11/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.sqrt(domain_guard((params['_0'] * x), min=0.1, nan=0.1))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * guarded_torch_power(params['_s'], torch.tensor(2)) * guarded_torch_power(x, torch.tensor(2))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_ones(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.sqrt(np_domain_guard((_0 * x), min=0.1, nan=0.1))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np_guarded_power(_s, np.array(2)) * np_guarded_power(x, np.array(2))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/11/loss.png b/fn_gen/ones_t/11/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..b364499e804c7a69ccacd86b3e8541d76faab338 Binary files /dev/null and b/fn_gen/ones_t/11/loss.png differ diff --git a/fn_gen/ones_t/11/quantization.png b/fn_gen/ones_t/11/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..e2ff36182f45ca4dec1e693e2b6aa826782c0ebf Binary files /dev/null and b/fn_gen/ones_t/11/quantization.png differ diff --git a/fn_gen/ones_t/12/distortion.png b/fn_gen/ones_t/12/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..52a65e02dd784e044a18ed54e045578a347d38de Binary files /dev/null and b/fn_gen/ones_t/12/distortion.png differ diff --git a/fn_gen/ones_t/12/expressions.txt b/fn_gen/ones_t/12/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..b835531ccc3a3813012a9a9487415f4f73afabc7 --- /dev/null +++ b/fn_gen/ones_t/12/expressions.txt @@ -0,0 +1,2 @@ +sinh(_0*x)/_s +log(_s*x - sqrt(_s**2*x**2 + 1))/_0 \ No newline at end of file diff --git a/fn_gen/ones_t/12/fn.py b/fn_gen/ones_t/12/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..252cef44f4819cd859f04a5000f7ce3086e549dd --- /dev/null +++ b/fn_gen/ones_t/12/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.sinh((params['_0'] * x))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.log(domain_guard(((torch.tensor(-1) * torch.sqrt(domain_guard((torch.tensor(1) + (guarded_torch_power(params['_s'], torch.tensor(2)) * guarded_torch_power(x, torch.tensor(2)))), min=0.1, nan=0.1))) + (params['_s'] * x)), min=1e-5, nan=1e-5))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_ones(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.sinh((_0 * x))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.log(np_domain_guard(((np.array(-1) * np.sqrt(np_domain_guard((np.array(1) + (np_guarded_power(_s, np.array(2)) * np_guarded_power(x, np.array(2)))), min=0.1, nan=0.1))) + (_s * x)), min=1e-5, nan=1e-5))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/12/loss.png b/fn_gen/ones_t/12/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..edf7f8937cea30056dbcfc49512d1dcc45873d36 Binary files /dev/null and b/fn_gen/ones_t/12/loss.png differ diff --git a/fn_gen/ones_t/12/quantization.png b/fn_gen/ones_t/12/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..8b40987077e138e7fb46b21e478139cfffb83ffb Binary files /dev/null and b/fn_gen/ones_t/12/quantization.png differ diff --git a/fn_gen/ones_t/13/distortion.png b/fn_gen/ones_t/13/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..9d387de9eb7f8643278bd71766853e856e34f91d Binary files /dev/null and b/fn_gen/ones_t/13/distortion.png differ diff --git a/fn_gen/ones_t/13/expressions.txt b/fn_gen/ones_t/13/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..6d6553d091cd1d343d7aa9b52b85ef6ec88ea854 --- /dev/null +++ b/fn_gen/ones_t/13/expressions.txt @@ -0,0 +1,2 @@ +x/_s +_s*x \ No newline at end of file diff --git a/fn_gen/ones_t/13/fn.py b/fn_gen/ones_t/13/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..059d0879691937eedbbd15f4c077a3720348bdf5 --- /dev/null +++ b/fn_gen/ones_t/13/fn.py @@ -0,0 +1,505 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (x * torch.div(1, replace_num(params['_s'], num=0, to=10000))) + + +def dequantization(x, **params): + return (params['_s'] * x) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _s): + return (x * np.divide(1, np_replace_num(_s, num=0, to=10000))) + + +def np_dequantization(x, _s): + return (_s * x) + + +def fit_func(x, _s): + x_ = np_quantization(x, _s) + x_ = np_dequantization(x_, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/13/loss.png b/fn_gen/ones_t/13/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..c386af1a6d4fa553268feb4e77c367f6ae0ed3f1 Binary files /dev/null and b/fn_gen/ones_t/13/loss.png differ diff --git a/fn_gen/ones_t/13/quantization.png b/fn_gen/ones_t/13/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..1079d2228f2452958bcf52f1dbc6ea4cc3444e99 Binary files /dev/null and b/fn_gen/ones_t/13/quantization.png differ diff --git a/fn_gen/ones_t/14/distortion.png b/fn_gen/ones_t/14/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..4d8aa0458769bdb35ef222c477970f089dacf9cb Binary files /dev/null and b/fn_gen/ones_t/14/distortion.png differ diff --git a/fn_gen/ones_t/14/expressions.txt b/fn_gen/ones_t/14/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..ed99293c42843616c361d59b23d32ae553cc0f8d --- /dev/null +++ b/fn_gen/ones_t/14/expressions.txt @@ -0,0 +1,2 @@ +atanh(_0*x)/_s +tanh(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/ones_t/14/fn.py b/fn_gen/ones_t/14/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..cf7e41383daa755f62bdda57703fb4622bc4dc66 --- /dev/null +++ b/fn_gen/ones_t/14/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.atanh(domain_guard((params['_0'] * x), min=-0.9999, max=0.9999, nan=0))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.tanh((params['_s'] * x))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_ones(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.arctanh(np_domain_guard((_0 * x), min=-0.9999, max=0.9999, nan=0))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.tanh((_s * x))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/14/loss.png b/fn_gen/ones_t/14/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..baefba0cd640d53d49aaf2caca64afd964ac4861 Binary files /dev/null and b/fn_gen/ones_t/14/loss.png differ diff --git a/fn_gen/ones_t/14/quantization.png b/fn_gen/ones_t/14/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..5717e6a9ed910576e9a35c2d5ff394d66d555b8a Binary files /dev/null and b/fn_gen/ones_t/14/quantization.png differ diff --git a/fn_gen/ones_t/15/distortion.png b/fn_gen/ones_t/15/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..0eedea0f1d1336b47f363fc3001e249d5663877d Binary files /dev/null and b/fn_gen/ones_t/15/distortion.png differ diff --git a/fn_gen/ones_t/15/expressions.txt b/fn_gen/ones_t/15/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..74791fc40576643d62f6366a8b4eda20eb1ad252 --- /dev/null +++ b/fn_gen/ones_t/15/expressions.txt @@ -0,0 +1,2 @@ +x**3/_s +(_s*x)**(1/3) \ No newline at end of file diff --git a/fn_gen/ones_t/15/fn.py b/fn_gen/ones_t/15/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..ec0681ddbf36f5c2b1f52a04072c9273576407be --- /dev/null +++ b/fn_gen/ones_t/15/fn.py @@ -0,0 +1,505 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * guarded_torch_power(x, torch.tensor(3))) + + +def dequantization(x, **params): + return guarded_torch_power((params['_s'] * x), 1 / 3) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np_guarded_power(x, np.array(3))) + + +def np_dequantization(x, _s): + return np_guarded_power((_s * x), 1 / 3) + + +def fit_func(x, _s): + x_ = np_quantization(x, _s) + x_ = np_dequantization(x_, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/15/loss.png b/fn_gen/ones_t/15/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..7fc9f03dded3fa2c095ce2846a6e14478915bc6c Binary files /dev/null and b/fn_gen/ones_t/15/loss.png differ diff --git a/fn_gen/ones_t/15/quantization.png b/fn_gen/ones_t/15/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..904ff18553babc0b9f6f96d41bb4dca0ee1ac169 Binary files /dev/null and b/fn_gen/ones_t/15/quantization.png differ diff --git a/fn_gen/ones_t/16/distortion.png b/fn_gen/ones_t/16/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..1985e0963da55e071b43fcd923d5b4afb6c8c91d Binary files /dev/null and b/fn_gen/ones_t/16/distortion.png differ diff --git a/fn_gen/ones_t/16/expressions.txt b/fn_gen/ones_t/16/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..23606e9f370f2e4adb43ed623c49d7fcaabd7355 --- /dev/null +++ b/fn_gen/ones_t/16/expressions.txt @@ -0,0 +1,2 @@ +tan(_0*x)/_s +atan(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/ones_t/16/fn.py b/fn_gen/ones_t/16/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..d6098a54347f634914905027130928e7a25e940c --- /dev/null +++ b/fn_gen/ones_t/16/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.tan(domain_guard((params['_0'] * x), posinf=1, neginf=-1, nan=0))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.atan((params['_s'] * x))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_ones(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.tan(np_domain_guard((_0 * x), posinf=1, neginf=-1, nan=0))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.arctan((_s * x))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/16/loss.png b/fn_gen/ones_t/16/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..21eebda07e761f1a60d319894825fbe06573fd0b Binary files /dev/null and b/fn_gen/ones_t/16/loss.png differ diff --git a/fn_gen/ones_t/16/quantization.png b/fn_gen/ones_t/16/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..6add8cdfd3d1e7f3dc7111bc345a0ac53739c39f Binary files /dev/null and b/fn_gen/ones_t/16/quantization.png differ diff --git a/fn_gen/ones_t/17/distortion.png b/fn_gen/ones_t/17/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..cfa28387318d24b05ae4a473a448e997be69d5e5 Binary files /dev/null and b/fn_gen/ones_t/17/distortion.png differ diff --git a/fn_gen/ones_t/17/expressions.txt b/fn_gen/ones_t/17/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..ec55493201f7b2b8effaefed75e0a9258fc25c56 --- /dev/null +++ b/fn_gen/ones_t/17/expressions.txt @@ -0,0 +1,2 @@ +tanh(_0*x)/_s +log((-_s*x - 1)/(_s*x - 1))/_0 \ No newline at end of file diff --git a/fn_gen/ones_t/17/fn.py b/fn_gen/ones_t/17/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..d1c49e5d2e30c9c5895f910a7a29896b85de5490 --- /dev/null +++ b/fn_gen/ones_t/17/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.tanh((params['_0'] * x))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.log(domain_guard((torch.div(1, replace_num((torch.tensor(-1) + (params['_s'] * x)), num=0, to=10000)) * (torch.tensor(-1) + (torch.tensor(-1) * params['_s'] * x))), min=1e-5, nan=1e-5))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_ones(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.tanh((_0 * x))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.log(np_domain_guard((np.divide(1, np_replace_num((np.array(-1) + (_s * x)), num=0, to=10000)) * (np.array(-1) + (np.array(-1) * _s * x))), min=1e-5, nan=1e-5))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/17/loss.png b/fn_gen/ones_t/17/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..a6111554de8f539f9f6262fc73e4ca8a5080f140 Binary files /dev/null and b/fn_gen/ones_t/17/loss.png differ diff --git a/fn_gen/ones_t/17/quantization.png b/fn_gen/ones_t/17/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..db1bf73f4453fe6c85be2d9c56b9d20d3843ae86 Binary files /dev/null and b/fn_gen/ones_t/17/quantization.png differ diff --git a/fn_gen/ones_t/18/distortion.png b/fn_gen/ones_t/18/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..f5e0fe5b12397e293c17c273a8b943abf7c146f0 Binary files /dev/null and b/fn_gen/ones_t/18/distortion.png differ diff --git a/fn_gen/ones_t/18/expressions.txt b/fn_gen/ones_t/18/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..03413827fa8f4c8ad49a40b543460cf31d1ce803 --- /dev/null +++ b/fn_gen/ones_t/18/expressions.txt @@ -0,0 +1,2 @@ +asin(_0*x)/_s +sin(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/ones_t/18/fn.py b/fn_gen/ones_t/18/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..213cccd4fc4e34eb5f2e98fadb71f1824be656e7 --- /dev/null +++ b/fn_gen/ones_t/18/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.asin(domain_guard((params['_0'] * x), min=-0.99999, max=0.99999, nan=0))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.sin((params['_s'] * x))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_ones(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.arcsin(np_domain_guard((_0 * x), min=-0.99999, max=0.99999, nan=0))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.sin((_s * x))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/18/loss.png b/fn_gen/ones_t/18/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..fdfd07737967e136df6578c90d004a0a2f68ea18 Binary files /dev/null and b/fn_gen/ones_t/18/loss.png differ diff --git a/fn_gen/ones_t/18/quantization.png b/fn_gen/ones_t/18/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..353c9d7a76da2df47f0cc3d89dcf4ddf26e9857b Binary files /dev/null and b/fn_gen/ones_t/18/quantization.png differ diff --git a/fn_gen/ones_t/2/distortion.png b/fn_gen/ones_t/2/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..f6f3ae86a9a8859cd5fe2ed07a3d3d4a8cca1a80 Binary files /dev/null and b/fn_gen/ones_t/2/distortion.png differ diff --git a/fn_gen/ones_t/2/expressions.txt b/fn_gen/ones_t/2/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..8c0b1579c06c048d5603aa39c80e392c5906a879 --- /dev/null +++ b/fn_gen/ones_t/2/expressions.txt @@ -0,0 +1,2 @@ +cos(_0*x)/_s +acos(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/ones_t/2/fn.py b/fn_gen/ones_t/2/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..88569b2b09f049cb704a87539cc1e22ff26e50fc --- /dev/null +++ b/fn_gen/ones_t/2/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.cos((params['_0'] * x))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.acos(domain_guard((params['_s'] * x), min=-0.99999, max=0.99999, nan=0))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_ones(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.cos((_0 * x))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.arccos(np_domain_guard((_s * x), min=-0.99999, max=0.99999, nan=0))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/2/loss.png b/fn_gen/ones_t/2/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..39dd0b0f2758ea480e4b85b224a5d7025bb07cdc Binary files /dev/null and b/fn_gen/ones_t/2/loss.png differ diff --git a/fn_gen/ones_t/2/quantization.png b/fn_gen/ones_t/2/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..42f03861f3d1ffe5f6c289bcef0149bc98f0bcef Binary files /dev/null and b/fn_gen/ones_t/2/quantization.png differ diff --git a/fn_gen/ones_t/3/distortion.png b/fn_gen/ones_t/3/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..a2ee371b1e8485c42a4c71d82570d92f84b5ccd0 Binary files /dev/null and b/fn_gen/ones_t/3/distortion.png differ diff --git a/fn_gen/ones_t/3/expressions.txt b/fn_gen/ones_t/3/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..dbb6da0fc54c6f23dc12daf2e2c3a395819e1bf4 --- /dev/null +++ b/fn_gen/ones_t/3/expressions.txt @@ -0,0 +1,2 @@ +x**2/_s +sqrt(_s*x) \ No newline at end of file diff --git a/fn_gen/ones_t/3/fn.py b/fn_gen/ones_t/3/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..a948ab3c71e46bbaa99a97d2dc5778b72e5a04d5 --- /dev/null +++ b/fn_gen/ones_t/3/fn.py @@ -0,0 +1,505 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * guarded_torch_power(x, torch.tensor(2))) + + +def dequantization(x, **params): + return torch.sqrt(domain_guard((params['_s'] * x), min=0.1, nan=0.1)) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np_guarded_power(x, np.array(2))) + + +def np_dequantization(x, _s): + return np.sqrt(np_domain_guard((_s * x), min=0.1, nan=0.1)) + + +def fit_func(x, _s): + x_ = np_quantization(x, _s) + x_ = np_dequantization(x_, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/3/loss.png b/fn_gen/ones_t/3/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..53821576df5b54c84fca2e60421abe311bf7ce61 Binary files /dev/null and b/fn_gen/ones_t/3/loss.png differ diff --git a/fn_gen/ones_t/3/quantization.png b/fn_gen/ones_t/3/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..114c44936bf8e8e42acc1426bb6a9100f9f30d25 Binary files /dev/null and b/fn_gen/ones_t/3/quantization.png differ diff --git a/fn_gen/ones_t/4/distortion.png b/fn_gen/ones_t/4/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..d78efc514dd75143064529b25295960e388a899d Binary files /dev/null and b/fn_gen/ones_t/4/distortion.png differ diff --git a/fn_gen/ones_t/4/expressions.txt b/fn_gen/ones_t/4/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..ecd6e238827dcdb95f4bcb390c1c300696f34254 --- /dev/null +++ b/fn_gen/ones_t/4/expressions.txt @@ -0,0 +1,2 @@ +sin(_0*x)/_s +asin(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/ones_t/4/fn.py b/fn_gen/ones_t/4/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..ca2cfc3da9ce76cc5f95159be9d96caebdc39929 --- /dev/null +++ b/fn_gen/ones_t/4/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.sin((params['_0'] * x))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.asin(domain_guard((params['_s'] * x), min=-0.99999, max=0.99999, nan=0))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_ones(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.sin((_0 * x))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.arcsin(np_domain_guard((_s * x), min=-0.99999, max=0.99999, nan=0))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/4/loss.png b/fn_gen/ones_t/4/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..fe14b0ee07ea8ce163dba81f5cc05a9ba57d0556 Binary files /dev/null and b/fn_gen/ones_t/4/loss.png differ diff --git a/fn_gen/ones_t/4/quantization.png b/fn_gen/ones_t/4/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..ecf0a9452e9d9dd50a590d8785a800470a8f5922 Binary files /dev/null and b/fn_gen/ones_t/4/quantization.png differ diff --git a/fn_gen/ones_t/5/distortion.png b/fn_gen/ones_t/5/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..656432a389fc8a12ef716b5aa5af064ea29b1178 Binary files /dev/null and b/fn_gen/ones_t/5/distortion.png differ diff --git a/fn_gen/ones_t/5/expressions.txt b/fn_gen/ones_t/5/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..9aa25379a9d1d5a93d60659c6609b2e24e79234d --- /dev/null +++ b/fn_gen/ones_t/5/expressions.txt @@ -0,0 +1,2 @@ +exp(_0*x)/_s +log(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/ones_t/5/fn.py b/fn_gen/ones_t/5/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..6271629e22caa00d86df9f0a74e64c054d42ae6e --- /dev/null +++ b/fn_gen/ones_t/5/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.exp((params['_0'] * x))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.log(domain_guard((params['_s'] * x), min=1e-5, nan=1e-5))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_ones(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.exp((_0 * x))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.log(np_domain_guard((_s * x), min=1e-5, nan=1e-5))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/5/loss.png b/fn_gen/ones_t/5/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..ccf51114d21dc61e81ca7e2175a49273a2a1d21c Binary files /dev/null and b/fn_gen/ones_t/5/loss.png differ diff --git a/fn_gen/ones_t/5/quantization.png b/fn_gen/ones_t/5/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..7febcf672e4b63a7226c94508243255c11a37fbf Binary files /dev/null and b/fn_gen/ones_t/5/quantization.png differ diff --git a/fn_gen/ones_t/6/distortion.png b/fn_gen/ones_t/6/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..9079821425d97882a154cd771317bfa9b68bc7cd Binary files /dev/null and b/fn_gen/ones_t/6/distortion.png differ diff --git a/fn_gen/ones_t/6/expressions.txt b/fn_gen/ones_t/6/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..5a7abbbdac98c7d53123fe0b9807e7644bc00acf --- /dev/null +++ b/fn_gen/ones_t/6/expressions.txt @@ -0,0 +1,2 @@ +acosh(_0*x)/_s +cosh(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/ones_t/6/fn.py b/fn_gen/ones_t/6/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..4fc8fb47a884d087f3d0da96d2a8bd22fa9adcfa --- /dev/null +++ b/fn_gen/ones_t/6/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.acosh(domain_guard((params['_0'] * x), min=1, nan=1))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.cosh((params['_s'] * x))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_ones(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.arccosh(np_domain_guard((_0 * x), min=1, nan=1))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.cosh((_s * x))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/6/loss.png b/fn_gen/ones_t/6/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..32610d21a344b65505b241efe34197786b533659 Binary files /dev/null and b/fn_gen/ones_t/6/loss.png differ diff --git a/fn_gen/ones_t/6/quantization.png b/fn_gen/ones_t/6/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..14c2e2f410a2f4b55c06ad5e1d223d7c4e054aee Binary files /dev/null and b/fn_gen/ones_t/6/quantization.png differ diff --git a/fn_gen/ones_t/7/distortion.png b/fn_gen/ones_t/7/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..ec375bee8de9b4f16e20d2f919cfe2d320f899ef Binary files /dev/null and b/fn_gen/ones_t/7/distortion.png differ diff --git a/fn_gen/ones_t/7/expressions.txt b/fn_gen/ones_t/7/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..c545adce8b3c320e195336b81461c79d0cc385e6 --- /dev/null +++ b/fn_gen/ones_t/7/expressions.txt @@ -0,0 +1,2 @@ +asinh(_0*x)/_s +sinh(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/ones_t/7/fn.py b/fn_gen/ones_t/7/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..09e843f5001bbce6c88a7a3bad7cec6f9c33a292 --- /dev/null +++ b/fn_gen/ones_t/7/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.asinh((params['_0'] * x))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.sinh((params['_s'] * x))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_ones(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.arcsinh((_0 * x))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.sinh((_s * x))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/7/loss.png b/fn_gen/ones_t/7/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..12d5a28346198f6fa4cc226a99ff0b9d2767643c Binary files /dev/null and b/fn_gen/ones_t/7/loss.png differ diff --git a/fn_gen/ones_t/7/quantization.png b/fn_gen/ones_t/7/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..0a132853ea18ce114f8a44d51fbd98d476e672eb Binary files /dev/null and b/fn_gen/ones_t/7/quantization.png differ diff --git a/fn_gen/ones_t/8/distortion.png b/fn_gen/ones_t/8/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..5daad1ce7ab02ebdf41b155f407e7dafdafc8da1 Binary files /dev/null and b/fn_gen/ones_t/8/distortion.png differ diff --git a/fn_gen/ones_t/8/expressions.txt b/fn_gen/ones_t/8/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..0a7e5be4566beeb4727d82f95d24241966d158dc --- /dev/null +++ b/fn_gen/ones_t/8/expressions.txt @@ -0,0 +1,2 @@ +log(_0*x)/_s +exp(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/ones_t/8/fn.py b/fn_gen/ones_t/8/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..b0f666064b1b3ad70c7e4bfa7310d4333954202d --- /dev/null +++ b/fn_gen/ones_t/8/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.log(domain_guard((params['_0'] * x), min=1e-5, nan=1e-5))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.exp((params['_s'] * x))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_ones(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.log(np_domain_guard((_0 * x), min=1e-5, nan=1e-5))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.exp((_s * x))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/8/loss.png b/fn_gen/ones_t/8/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..37c0f6825ece180b73abfe667cacd973d082ea40 Binary files /dev/null and b/fn_gen/ones_t/8/loss.png differ diff --git a/fn_gen/ones_t/8/quantization.png b/fn_gen/ones_t/8/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..278cb5e8d1f383e7262a71538c6fdb551d33c169 Binary files /dev/null and b/fn_gen/ones_t/8/quantization.png differ diff --git a/fn_gen/ones_t/9/distortion.png b/fn_gen/ones_t/9/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..b0c0444cdac99007f990b893bd0bdfd08a32005f Binary files /dev/null and b/fn_gen/ones_t/9/distortion.png differ diff --git a/fn_gen/ones_t/9/expressions.txt b/fn_gen/ones_t/9/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..aa32b575e8c654dbc457c94f36222e70d86dc940 --- /dev/null +++ b/fn_gen/ones_t/9/expressions.txt @@ -0,0 +1,2 @@ +atan(_0*x)/_s +tan(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/ones_t/9/fn.py b/fn_gen/ones_t/9/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..bdcb8a8ab9dfe6384d478403ccd9d99d5ae8c464 --- /dev/null +++ b/fn_gen/ones_t/9/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.atan((params['_0'] * x))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.tan(domain_guard((params['_s'] * x), posinf=1, neginf=-1, nan=0))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_ones(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.arctan((_0 * x))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.tan(np_domain_guard((_s * x), posinf=1, neginf=-1, nan=0))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/ones_t/9/loss.png b/fn_gen/ones_t/9/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..a4f5f26d134547d388b962455a946487ff5810f3 Binary files /dev/null and b/fn_gen/ones_t/9/loss.png differ diff --git a/fn_gen/ones_t/9/quantization.png b/fn_gen/ones_t/9/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..94eea6e7af5150c70fb258e2f41c1a9067be7e68 Binary files /dev/null and b/fn_gen/ones_t/9/quantization.png differ diff --git a/fn_gen/rnd_naive_t/0/distortion.png b/fn_gen/rnd_naive_t/0/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..eb5563f2d9704331a8db88886be7693fedcdc72d Binary files /dev/null and b/fn_gen/rnd_naive_t/0/distortion.png differ diff --git a/fn_gen/rnd_naive_t/0/expressions.txt b/fn_gen/rnd_naive_t/0/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..3758ee2a62aa8d95c3b7da1dd3fafa11b027ad9b --- /dev/null +++ b/fn_gen/rnd_naive_t/0/expressions.txt @@ -0,0 +1,2 @@ +cosh(_0*x)/_s +log(_s*x - sqrt(_s**2*x**2 - 1))/_0 \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/0/fn.py b/fn_gen/rnd_naive_t/0/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..3072f861c7468b079aa8defaab6a5caddd94c2e4 --- /dev/null +++ b/fn_gen/rnd_naive_t/0/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.cosh((params['_0'] * x))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.log(domain_guard(((torch.tensor(-1) * torch.sqrt(domain_guard((torch.tensor(-1) + (guarded_torch_power(params['_s'], torch.tensor(2)) * guarded_torch_power(x, torch.tensor(2)))), min=0.1, nan=0.1))) + (params['_s'] * x)), min=1e-5, nan=1e-5))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_rand(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.cosh((_0 * x))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.log(np_domain_guard(((np.array(-1) * np.sqrt(np_domain_guard((np.array(-1) + (np_guarded_power(_s, np.array(2)) * np_guarded_power(x, np.array(2)))), min=0.1, nan=0.1))) + (_s * x)), min=1e-5, nan=1e-5))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/0/loss.png b/fn_gen/rnd_naive_t/0/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..9b202758d3ce3b68bd1e4163dd6d6c72f8d86f98 Binary files /dev/null and b/fn_gen/rnd_naive_t/0/loss.png differ diff --git a/fn_gen/rnd_naive_t/0/quantization.png b/fn_gen/rnd_naive_t/0/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..31cd59437026e1f27b9eec5b89570f6e060ed276 Binary files /dev/null and b/fn_gen/rnd_naive_t/0/quantization.png differ diff --git a/fn_gen/rnd_naive_t/1/distortion.png b/fn_gen/rnd_naive_t/1/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..59df2ad696eb3c2cbae28407531ed016a314e4ef Binary files /dev/null and b/fn_gen/rnd_naive_t/1/distortion.png differ diff --git a/fn_gen/rnd_naive_t/1/expressions.txt b/fn_gen/rnd_naive_t/1/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..03413827fa8f4c8ad49a40b543460cf31d1ce803 --- /dev/null +++ b/fn_gen/rnd_naive_t/1/expressions.txt @@ -0,0 +1,2 @@ +asin(_0*x)/_s +sin(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/1/fn.py b/fn_gen/rnd_naive_t/1/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..2f789a24dc19cb3a6e3b83209ab97d4513b98352 --- /dev/null +++ b/fn_gen/rnd_naive_t/1/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.asin(domain_guard((params['_0'] * x), min=-0.99999, max=0.99999, nan=0))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.sin((params['_s'] * x))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_rand(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.arcsin(np_domain_guard((_0 * x), min=-0.99999, max=0.99999, nan=0))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.sin((_s * x))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/1/loss.png b/fn_gen/rnd_naive_t/1/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..8b0392488f06d3c8b642799a7e9b6e5aca02c1ef Binary files /dev/null and b/fn_gen/rnd_naive_t/1/loss.png differ diff --git a/fn_gen/rnd_naive_t/1/quantization.png b/fn_gen/rnd_naive_t/1/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..8bd871949c794719ec131943c38be3fc6bafc8b0 Binary files /dev/null and b/fn_gen/rnd_naive_t/1/quantization.png differ diff --git a/fn_gen/rnd_naive_t/10/distortion.png b/fn_gen/rnd_naive_t/10/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..ccb6e2e0173d7fa41677e0fb8571906ab584eab1 Binary files /dev/null and b/fn_gen/rnd_naive_t/10/distortion.png differ diff --git a/fn_gen/rnd_naive_t/10/expressions.txt b/fn_gen/rnd_naive_t/10/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..e8458af52eb4cfce21cf8459f3c454003cd78158 --- /dev/null +++ b/fn_gen/rnd_naive_t/10/expressions.txt @@ -0,0 +1,2 @@ +sqrt(_0*x)/_s +_s**2*x**2/_0 \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/10/fn.py b/fn_gen/rnd_naive_t/10/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..b260614002c72042073dbd913a5ec07f464d1eee --- /dev/null +++ b/fn_gen/rnd_naive_t/10/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.sqrt(domain_guard((params['_0'] * x), min=0.1, nan=0.1))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * guarded_torch_power(params['_s'], torch.tensor(2)) * guarded_torch_power(x, torch.tensor(2))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_rand(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.sqrt(np_domain_guard((_0 * x), min=0.1, nan=0.1))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np_guarded_power(_s, np.array(2)) * np_guarded_power(x, np.array(2))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/10/loss.png b/fn_gen/rnd_naive_t/10/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..fc079c433c722ec14f8087355406e542fed5bb87 Binary files /dev/null and b/fn_gen/rnd_naive_t/10/loss.png differ diff --git a/fn_gen/rnd_naive_t/10/quantization.png b/fn_gen/rnd_naive_t/10/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..af1bdec83711bba0e6a940dc4b21f5e03cc1b249 Binary files /dev/null and b/fn_gen/rnd_naive_t/10/quantization.png differ diff --git a/fn_gen/rnd_naive_t/11/distortion.png b/fn_gen/rnd_naive_t/11/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..54e13423c3a11786e8bebb5010441c29808fa8d9 Binary files /dev/null and b/fn_gen/rnd_naive_t/11/distortion.png differ diff --git a/fn_gen/rnd_naive_t/11/expressions.txt b/fn_gen/rnd_naive_t/11/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..ed99293c42843616c361d59b23d32ae553cc0f8d --- /dev/null +++ b/fn_gen/rnd_naive_t/11/expressions.txt @@ -0,0 +1,2 @@ +atanh(_0*x)/_s +tanh(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/11/fn.py b/fn_gen/rnd_naive_t/11/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..7b7a57dd104c6f01bc374d45b7f1b52fecf8fa9a --- /dev/null +++ b/fn_gen/rnd_naive_t/11/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.atanh(domain_guard((params['_0'] * x), min=-0.9999, max=0.9999, nan=0))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.tanh((params['_s'] * x))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_rand(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.arctanh(np_domain_guard((_0 * x), min=-0.9999, max=0.9999, nan=0))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.tanh((_s * x))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/11/loss.png b/fn_gen/rnd_naive_t/11/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..b566995e8eaf5100d658b15ebda67867ea5c5900 Binary files /dev/null and b/fn_gen/rnd_naive_t/11/loss.png differ diff --git a/fn_gen/rnd_naive_t/11/quantization.png b/fn_gen/rnd_naive_t/11/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..3ecfcc5c6b5701cf8cfbdef47f0ab374456dc8a4 Binary files /dev/null and b/fn_gen/rnd_naive_t/11/quantization.png differ diff --git a/fn_gen/rnd_naive_t/12/distortion.png b/fn_gen/rnd_naive_t/12/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..b1e97af290065b1f240d856b851cda969b7eb749 Binary files /dev/null and b/fn_gen/rnd_naive_t/12/distortion.png differ diff --git a/fn_gen/rnd_naive_t/12/expressions.txt b/fn_gen/rnd_naive_t/12/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..b835531ccc3a3813012a9a9487415f4f73afabc7 --- /dev/null +++ b/fn_gen/rnd_naive_t/12/expressions.txt @@ -0,0 +1,2 @@ +sinh(_0*x)/_s +log(_s*x - sqrt(_s**2*x**2 + 1))/_0 \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/12/fn.py b/fn_gen/rnd_naive_t/12/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..2a4543980b6e64b1a6e75871fde0abbdb44ea5fd --- /dev/null +++ b/fn_gen/rnd_naive_t/12/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.sinh((params['_0'] * x))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.log(domain_guard(((torch.tensor(-1) * torch.sqrt(domain_guard((torch.tensor(1) + (guarded_torch_power(params['_s'], torch.tensor(2)) * guarded_torch_power(x, torch.tensor(2)))), min=0.1, nan=0.1))) + (params['_s'] * x)), min=1e-5, nan=1e-5))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_rand(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.sinh((_0 * x))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.log(np_domain_guard(((np.array(-1) * np.sqrt(np_domain_guard((np.array(1) + (np_guarded_power(_s, np.array(2)) * np_guarded_power(x, np.array(2)))), min=0.1, nan=0.1))) + (_s * x)), min=1e-5, nan=1e-5))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/12/loss.png b/fn_gen/rnd_naive_t/12/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..1ce8c6294cd8bf923db9350b97ff88755fb5b273 Binary files /dev/null and b/fn_gen/rnd_naive_t/12/loss.png differ diff --git a/fn_gen/rnd_naive_t/12/quantization.png b/fn_gen/rnd_naive_t/12/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..814b49c667eb2b11c0f574818341af99eebbe86f Binary files /dev/null and b/fn_gen/rnd_naive_t/12/quantization.png differ diff --git a/fn_gen/rnd_naive_t/13/distortion.png b/fn_gen/rnd_naive_t/13/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..176c125008c4ee88870008f2501f12b2853808f6 Binary files /dev/null and b/fn_gen/rnd_naive_t/13/distortion.png differ diff --git a/fn_gen/rnd_naive_t/13/expressions.txt b/fn_gen/rnd_naive_t/13/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..ec55493201f7b2b8effaefed75e0a9258fc25c56 --- /dev/null +++ b/fn_gen/rnd_naive_t/13/expressions.txt @@ -0,0 +1,2 @@ +tanh(_0*x)/_s +log((-_s*x - 1)/(_s*x - 1))/_0 \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/13/fn.py b/fn_gen/rnd_naive_t/13/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..fd4b07a9cfb29bfc3c6c28082be288fe156c5099 --- /dev/null +++ b/fn_gen/rnd_naive_t/13/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.tanh((params['_0'] * x))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.log(domain_guard((torch.div(1, replace_num((torch.tensor(-1) + (params['_s'] * x)), num=0, to=10000)) * (torch.tensor(-1) + (torch.tensor(-1) * params['_s'] * x))), min=1e-5, nan=1e-5))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_rand(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.tanh((_0 * x))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.log(np_domain_guard((np.divide(1, np_replace_num((np.array(-1) + (_s * x)), num=0, to=10000)) * (np.array(-1) + (np.array(-1) * _s * x))), min=1e-5, nan=1e-5))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/13/loss.png b/fn_gen/rnd_naive_t/13/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..8eea87b6ea43174eb28ba5fd8d092c60f6c60aeb Binary files /dev/null and b/fn_gen/rnd_naive_t/13/loss.png differ diff --git a/fn_gen/rnd_naive_t/13/quantization.png b/fn_gen/rnd_naive_t/13/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..36f1112d89194419020e41cbfcc610982f372948 Binary files /dev/null and b/fn_gen/rnd_naive_t/13/quantization.png differ diff --git a/fn_gen/rnd_naive_t/14/distortion.png b/fn_gen/rnd_naive_t/14/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..fa03e23de1605b80812dc75c9c408e2616414d62 Binary files /dev/null and b/fn_gen/rnd_naive_t/14/distortion.png differ diff --git a/fn_gen/rnd_naive_t/14/expressions.txt b/fn_gen/rnd_naive_t/14/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..c7b68c388fdf6e1b6e2be8076f1d4b8d7bcef4f9 --- /dev/null +++ b/fn_gen/rnd_naive_t/14/expressions.txt @@ -0,0 +1,2 @@ +(_0*x)**(1/3)/_s +_s**3*x**3/_0 \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/14/fn.py b/fn_gen/rnd_naive_t/14/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..7f045f15f3bfb8f7558b2b671a4b783137bb8ef4 --- /dev/null +++ b/fn_gen/rnd_naive_t/14/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * guarded_torch_power((params['_0'] * x), 1 / 3)) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * guarded_torch_power(params['_s'], torch.tensor(3)) * guarded_torch_power(x, torch.tensor(3))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_rand(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np_guarded_power((_0 * x), 1 / 3)) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np_guarded_power(_s, np.array(3)) * np_guarded_power(x, np.array(3))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/14/loss.png b/fn_gen/rnd_naive_t/14/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..124aaae27d7b3fb2c4022f57af26e7a006564eeb Binary files /dev/null and b/fn_gen/rnd_naive_t/14/loss.png differ diff --git a/fn_gen/rnd_naive_t/14/quantization.png b/fn_gen/rnd_naive_t/14/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..6b64d4ade8488c367b1f59c188411843ca4ea619 Binary files /dev/null and b/fn_gen/rnd_naive_t/14/quantization.png differ diff --git a/fn_gen/rnd_naive_t/15/distortion.png b/fn_gen/rnd_naive_t/15/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..9d387de9eb7f8643278bd71766853e856e34f91d Binary files /dev/null and b/fn_gen/rnd_naive_t/15/distortion.png differ diff --git a/fn_gen/rnd_naive_t/15/expressions.txt b/fn_gen/rnd_naive_t/15/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..6d6553d091cd1d343d7aa9b52b85ef6ec88ea854 --- /dev/null +++ b/fn_gen/rnd_naive_t/15/expressions.txt @@ -0,0 +1,2 @@ +x/_s +_s*x \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/15/fn.py b/fn_gen/rnd_naive_t/15/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..059d0879691937eedbbd15f4c077a3720348bdf5 --- /dev/null +++ b/fn_gen/rnd_naive_t/15/fn.py @@ -0,0 +1,505 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (x * torch.div(1, replace_num(params['_s'], num=0, to=10000))) + + +def dequantization(x, **params): + return (params['_s'] * x) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _s): + return (x * np.divide(1, np_replace_num(_s, num=0, to=10000))) + + +def np_dequantization(x, _s): + return (_s * x) + + +def fit_func(x, _s): + x_ = np_quantization(x, _s) + x_ = np_dequantization(x_, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/15/loss.png b/fn_gen/rnd_naive_t/15/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..73b4f010733d100afa33f3982861305854ffbb87 Binary files /dev/null and b/fn_gen/rnd_naive_t/15/loss.png differ diff --git a/fn_gen/rnd_naive_t/15/quantization.png b/fn_gen/rnd_naive_t/15/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..1079d2228f2452958bcf52f1dbc6ea4cc3444e99 Binary files /dev/null and b/fn_gen/rnd_naive_t/15/quantization.png differ diff --git a/fn_gen/rnd_naive_t/16/distortion.png b/fn_gen/rnd_naive_t/16/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..f4711e60d5e2e5d2745faeb587b6552245f804fc Binary files /dev/null and b/fn_gen/rnd_naive_t/16/distortion.png differ diff --git a/fn_gen/rnd_naive_t/16/expressions.txt b/fn_gen/rnd_naive_t/16/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..576ec6a351e26f9982eb17e394804ca906d4b067 --- /dev/null +++ b/fn_gen/rnd_naive_t/16/expressions.txt @@ -0,0 +1,2 @@ +acos(_0*x)/_s +cos(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/16/fn.py b/fn_gen/rnd_naive_t/16/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..114bc3c2e786c7c43cd4742cb1387c2cb84556b6 --- /dev/null +++ b/fn_gen/rnd_naive_t/16/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.acos(domain_guard((params['_0'] * x), min=-0.99999, max=0.99999, nan=0))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.cos((params['_s'] * x))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_rand(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.arccos(np_domain_guard((_0 * x), min=-0.99999, max=0.99999, nan=0))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.cos((_s * x))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/16/loss.png b/fn_gen/rnd_naive_t/16/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..399682d59478403a52c68f48e7ec6851f97ac413 Binary files /dev/null and b/fn_gen/rnd_naive_t/16/loss.png differ diff --git a/fn_gen/rnd_naive_t/16/quantization.png b/fn_gen/rnd_naive_t/16/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..5b774f61c2d8cde564bc53489b8aff3c8f3a2da4 Binary files /dev/null and b/fn_gen/rnd_naive_t/16/quantization.png differ diff --git a/fn_gen/rnd_naive_t/17/distortion.png b/fn_gen/rnd_naive_t/17/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..63fe12e897a8c1442b0efd34efaed7211acdc50e Binary files /dev/null and b/fn_gen/rnd_naive_t/17/distortion.png differ diff --git a/fn_gen/rnd_naive_t/17/expressions.txt b/fn_gen/rnd_naive_t/17/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..23606e9f370f2e4adb43ed623c49d7fcaabd7355 --- /dev/null +++ b/fn_gen/rnd_naive_t/17/expressions.txt @@ -0,0 +1,2 @@ +tan(_0*x)/_s +atan(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/17/fn.py b/fn_gen/rnd_naive_t/17/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..922ebfa319ae1df327f501ae1ce9251bdfc6c24b --- /dev/null +++ b/fn_gen/rnd_naive_t/17/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.tan(domain_guard((params['_0'] * x), posinf=1, neginf=-1, nan=0))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.atan((params['_s'] * x))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_rand(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.tan(np_domain_guard((_0 * x), posinf=1, neginf=-1, nan=0))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.arctan((_s * x))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/17/loss.png b/fn_gen/rnd_naive_t/17/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..6e869ffa60446d11a5b2c035c8b07d045109685e Binary files /dev/null and b/fn_gen/rnd_naive_t/17/loss.png differ diff --git a/fn_gen/rnd_naive_t/17/quantization.png b/fn_gen/rnd_naive_t/17/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..6d9195236768d7d71fe4326f6dd1b711c7e545c8 Binary files /dev/null and b/fn_gen/rnd_naive_t/17/quantization.png differ diff --git a/fn_gen/rnd_naive_t/18/distortion.png b/fn_gen/rnd_naive_t/18/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..12625a10c2d648940bfbc8c875774efe5bf5653c Binary files /dev/null and b/fn_gen/rnd_naive_t/18/distortion.png differ diff --git a/fn_gen/rnd_naive_t/18/expressions.txt b/fn_gen/rnd_naive_t/18/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..c545adce8b3c320e195336b81461c79d0cc385e6 --- /dev/null +++ b/fn_gen/rnd_naive_t/18/expressions.txt @@ -0,0 +1,2 @@ +asinh(_0*x)/_s +sinh(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/18/fn.py b/fn_gen/rnd_naive_t/18/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..af134fb8030e329e83799fb4c0f86ce19ab85308 --- /dev/null +++ b/fn_gen/rnd_naive_t/18/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.asinh((params['_0'] * x))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.sinh((params['_s'] * x))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_rand(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.arcsinh((_0 * x))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.sinh((_s * x))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/18/loss.png b/fn_gen/rnd_naive_t/18/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..7228af5fb8a7986d51d56bf3a83fdb117d5171f8 Binary files /dev/null and b/fn_gen/rnd_naive_t/18/loss.png differ diff --git a/fn_gen/rnd_naive_t/18/quantization.png b/fn_gen/rnd_naive_t/18/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..55c7f25fe4af7abc9a2ccea0a2bb460f86c6e349 Binary files /dev/null and b/fn_gen/rnd_naive_t/18/quantization.png differ diff --git a/fn_gen/rnd_naive_t/2/distortion.png b/fn_gen/rnd_naive_t/2/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..3b5e9888b4663eb1bce1e3423d10586de6b92ff5 Binary files /dev/null and b/fn_gen/rnd_naive_t/2/distortion.png differ diff --git a/fn_gen/rnd_naive_t/2/expressions.txt b/fn_gen/rnd_naive_t/2/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..9aa25379a9d1d5a93d60659c6609b2e24e79234d --- /dev/null +++ b/fn_gen/rnd_naive_t/2/expressions.txt @@ -0,0 +1,2 @@ +exp(_0*x)/_s +log(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/2/fn.py b/fn_gen/rnd_naive_t/2/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..5574a78cc00e708d624e259a1bfa6710252167e1 --- /dev/null +++ b/fn_gen/rnd_naive_t/2/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.exp((params['_0'] * x))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.log(domain_guard((params['_s'] * x), min=1e-5, nan=1e-5))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_rand(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.exp((_0 * x))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.log(np_domain_guard((_s * x), min=1e-5, nan=1e-5))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/2/loss.png b/fn_gen/rnd_naive_t/2/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..4da62064f88f825c8653a511c91c51319e0d00b5 Binary files /dev/null and b/fn_gen/rnd_naive_t/2/loss.png differ diff --git a/fn_gen/rnd_naive_t/2/quantization.png b/fn_gen/rnd_naive_t/2/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..ec12c07926948876b71e3b6650dbb4dc773c9531 Binary files /dev/null and b/fn_gen/rnd_naive_t/2/quantization.png differ diff --git a/fn_gen/rnd_naive_t/3/distortion.png b/fn_gen/rnd_naive_t/3/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..0eedea0f1d1336b47f363fc3001e249d5663877d Binary files /dev/null and b/fn_gen/rnd_naive_t/3/distortion.png differ diff --git a/fn_gen/rnd_naive_t/3/expressions.txt b/fn_gen/rnd_naive_t/3/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..74791fc40576643d62f6366a8b4eda20eb1ad252 --- /dev/null +++ b/fn_gen/rnd_naive_t/3/expressions.txt @@ -0,0 +1,2 @@ +x**3/_s +(_s*x)**(1/3) \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/3/fn.py b/fn_gen/rnd_naive_t/3/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..ec0681ddbf36f5c2b1f52a04072c9273576407be --- /dev/null +++ b/fn_gen/rnd_naive_t/3/fn.py @@ -0,0 +1,505 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * guarded_torch_power(x, torch.tensor(3))) + + +def dequantization(x, **params): + return guarded_torch_power((params['_s'] * x), 1 / 3) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np_guarded_power(x, np.array(3))) + + +def np_dequantization(x, _s): + return np_guarded_power((_s * x), 1 / 3) + + +def fit_func(x, _s): + x_ = np_quantization(x, _s) + x_ = np_dequantization(x_, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/3/loss.png b/fn_gen/rnd_naive_t/3/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..6dc3aa9856e573854f011b9564e05eee84478bfd Binary files /dev/null and b/fn_gen/rnd_naive_t/3/loss.png differ diff --git a/fn_gen/rnd_naive_t/3/quantization.png b/fn_gen/rnd_naive_t/3/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..904ff18553babc0b9f6f96d41bb4dca0ee1ac169 Binary files /dev/null and b/fn_gen/rnd_naive_t/3/quantization.png differ diff --git a/fn_gen/rnd_naive_t/4/distortion.png b/fn_gen/rnd_naive_t/4/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..550d6f34cf4a09899024c6306e3e425d262fe506 Binary files /dev/null and b/fn_gen/rnd_naive_t/4/distortion.png differ diff --git a/fn_gen/rnd_naive_t/4/expressions.txt b/fn_gen/rnd_naive_t/4/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..5a7abbbdac98c7d53123fe0b9807e7644bc00acf --- /dev/null +++ b/fn_gen/rnd_naive_t/4/expressions.txt @@ -0,0 +1,2 @@ +acosh(_0*x)/_s +cosh(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/4/fn.py b/fn_gen/rnd_naive_t/4/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..c7cd02b97c0fa312bb6af767d82eaec13bfd76c5 --- /dev/null +++ b/fn_gen/rnd_naive_t/4/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.acosh(domain_guard((params['_0'] * x), min=1, nan=1))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.cosh((params['_s'] * x))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_rand(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.arccosh(np_domain_guard((_0 * x), min=1, nan=1))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.cosh((_s * x))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/4/loss.png b/fn_gen/rnd_naive_t/4/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..87a60d01c14065298872875be65b01593fdb076c Binary files /dev/null and b/fn_gen/rnd_naive_t/4/loss.png differ diff --git a/fn_gen/rnd_naive_t/4/quantization.png b/fn_gen/rnd_naive_t/4/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..14c2e2f410a2f4b55c06ad5e1d223d7c4e054aee Binary files /dev/null and b/fn_gen/rnd_naive_t/4/quantization.png differ diff --git a/fn_gen/rnd_naive_t/5/distortion.png b/fn_gen/rnd_naive_t/5/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..ea40c0c6ced1c6a9f0e1c329050d60a64e4cdcc6 Binary files /dev/null and b/fn_gen/rnd_naive_t/5/distortion.png differ diff --git a/fn_gen/rnd_naive_t/5/expressions.txt b/fn_gen/rnd_naive_t/5/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..0a7e5be4566beeb4727d82f95d24241966d158dc --- /dev/null +++ b/fn_gen/rnd_naive_t/5/expressions.txt @@ -0,0 +1,2 @@ +log(_0*x)/_s +exp(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/5/fn.py b/fn_gen/rnd_naive_t/5/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..c4017b5c9ad111e483a5db84df5d9ae1d6bc20df --- /dev/null +++ b/fn_gen/rnd_naive_t/5/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.log(domain_guard((params['_0'] * x), min=1e-5, nan=1e-5))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.exp((params['_s'] * x))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_rand(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.log(np_domain_guard((_0 * x), min=1e-5, nan=1e-5))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.exp((_s * x))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/5/loss.png b/fn_gen/rnd_naive_t/5/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..cc9ee9dd9692ad0690aa68bc7b04983c54a9bc0f Binary files /dev/null and b/fn_gen/rnd_naive_t/5/loss.png differ diff --git a/fn_gen/rnd_naive_t/5/quantization.png b/fn_gen/rnd_naive_t/5/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..542701c3f8ec37a086748c69a6cbb602a3c3a57e Binary files /dev/null and b/fn_gen/rnd_naive_t/5/quantization.png differ diff --git a/fn_gen/rnd_naive_t/6/distortion.png b/fn_gen/rnd_naive_t/6/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..c1aca2d82d04ab3892028b0e8c82951e44e71fca Binary files /dev/null and b/fn_gen/rnd_naive_t/6/distortion.png differ diff --git a/fn_gen/rnd_naive_t/6/expressions.txt b/fn_gen/rnd_naive_t/6/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..aa32b575e8c654dbc457c94f36222e70d86dc940 --- /dev/null +++ b/fn_gen/rnd_naive_t/6/expressions.txt @@ -0,0 +1,2 @@ +atan(_0*x)/_s +tan(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/6/fn.py b/fn_gen/rnd_naive_t/6/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..e3aeb42e9b9a9e3b0ecec2c8b181cfcaa669c4a8 --- /dev/null +++ b/fn_gen/rnd_naive_t/6/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.atan((params['_0'] * x))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.tan(domain_guard((params['_s'] * x), posinf=1, neginf=-1, nan=0))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_rand(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.arctan((_0 * x))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.tan(np_domain_guard((_s * x), posinf=1, neginf=-1, nan=0))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/6/loss.png b/fn_gen/rnd_naive_t/6/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..39dbb909eb6769080daaf164d01e2fbc3d73c14d Binary files /dev/null and b/fn_gen/rnd_naive_t/6/loss.png differ diff --git a/fn_gen/rnd_naive_t/6/quantization.png b/fn_gen/rnd_naive_t/6/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..43668ad491bd28cc6711fc96b3aedf4951ca1607 Binary files /dev/null and b/fn_gen/rnd_naive_t/6/quantization.png differ diff --git a/fn_gen/rnd_naive_t/7/distortion.png b/fn_gen/rnd_naive_t/7/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..a2ee371b1e8485c42a4c71d82570d92f84b5ccd0 Binary files /dev/null and b/fn_gen/rnd_naive_t/7/distortion.png differ diff --git a/fn_gen/rnd_naive_t/7/expressions.txt b/fn_gen/rnd_naive_t/7/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..dbb6da0fc54c6f23dc12daf2e2c3a395819e1bf4 --- /dev/null +++ b/fn_gen/rnd_naive_t/7/expressions.txt @@ -0,0 +1,2 @@ +x**2/_s +sqrt(_s*x) \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/7/fn.py b/fn_gen/rnd_naive_t/7/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..a948ab3c71e46bbaa99a97d2dc5778b72e5a04d5 --- /dev/null +++ b/fn_gen/rnd_naive_t/7/fn.py @@ -0,0 +1,505 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * guarded_torch_power(x, torch.tensor(2))) + + +def dequantization(x, **params): + return torch.sqrt(domain_guard((params['_s'] * x), min=0.1, nan=0.1)) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np_guarded_power(x, np.array(2))) + + +def np_dequantization(x, _s): + return np.sqrt(np_domain_guard((_s * x), min=0.1, nan=0.1)) + + +def fit_func(x, _s): + x_ = np_quantization(x, _s) + x_ = np_dequantization(x_, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/7/loss.png b/fn_gen/rnd_naive_t/7/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..5b19aa22d2d6a4973eb18e8197cf025aed824988 Binary files /dev/null and b/fn_gen/rnd_naive_t/7/loss.png differ diff --git a/fn_gen/rnd_naive_t/7/quantization.png b/fn_gen/rnd_naive_t/7/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..114c44936bf8e8e42acc1426bb6a9100f9f30d25 Binary files /dev/null and b/fn_gen/rnd_naive_t/7/quantization.png differ diff --git a/fn_gen/rnd_naive_t/8/distortion.png b/fn_gen/rnd_naive_t/8/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..af6e614ff598636c9c7007da21f262fe11f2225d Binary files /dev/null and b/fn_gen/rnd_naive_t/8/distortion.png differ diff --git a/fn_gen/rnd_naive_t/8/expressions.txt b/fn_gen/rnd_naive_t/8/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..8c0b1579c06c048d5603aa39c80e392c5906a879 --- /dev/null +++ b/fn_gen/rnd_naive_t/8/expressions.txt @@ -0,0 +1,2 @@ +cos(_0*x)/_s +acos(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/8/fn.py b/fn_gen/rnd_naive_t/8/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..edfdc0940a01e2b9fea82b52232889212469f17f --- /dev/null +++ b/fn_gen/rnd_naive_t/8/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.cos((params['_0'] * x))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.acos(domain_guard((params['_s'] * x), min=-0.99999, max=0.99999, nan=0))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_rand(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.cos((_0 * x))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.arccos(np_domain_guard((_s * x), min=-0.99999, max=0.99999, nan=0))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/8/loss.png b/fn_gen/rnd_naive_t/8/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..8b67583e05d7458bec68df41b58f5d89ca2f6e91 Binary files /dev/null and b/fn_gen/rnd_naive_t/8/loss.png differ diff --git a/fn_gen/rnd_naive_t/8/quantization.png b/fn_gen/rnd_naive_t/8/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..072938142fe0fd73d72d3a13266ee91c303fc503 Binary files /dev/null and b/fn_gen/rnd_naive_t/8/quantization.png differ diff --git a/fn_gen/rnd_naive_t/9/distortion.png b/fn_gen/rnd_naive_t/9/distortion.png new file mode 100644 index 0000000000000000000000000000000000000000..7eded65215a0ec85271922204dd6d177b204326f Binary files /dev/null and b/fn_gen/rnd_naive_t/9/distortion.png differ diff --git a/fn_gen/rnd_naive_t/9/expressions.txt b/fn_gen/rnd_naive_t/9/expressions.txt new file mode 100644 index 0000000000000000000000000000000000000000..ecd6e238827dcdb95f4bcb390c1c300696f34254 --- /dev/null +++ b/fn_gen/rnd_naive_t/9/expressions.txt @@ -0,0 +1,2 @@ +sin(_0*x)/_s +asin(_s*x)/_0 \ No newline at end of file diff --git a/fn_gen/rnd_naive_t/9/fn.py b/fn_gen/rnd_naive_t/9/fn.py new file mode 100644 index 0000000000000000000000000000000000000000..c1d2a7b4905165147008a44c1f9ed00dbe993606 --- /dev/null +++ b/fn_gen/rnd_naive_t/9/fn.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import torch +from torch import amin # Necessary for arcsin +import copy +import torch.nn as nn +import numpy as np + +from scipy.optimize import curve_fit +from typing import Dict, Any, Tuple, List, Callable + + +def quantization(x, **params): + return (torch.div(1, replace_num(params['_s'], num=0, to=10000)) * torch.sin((params['_0'] * x))) + + +def dequantization(x, **params): + return (torch.div(1, replace_num(params['_0'], num=0, to=10000)) * torch.asin(domain_guard((params['_s'] * x), min=-0.99999, max=0.99999, nan=0))) + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_params(x: torch.Tensor, **kwargs: Dict[str, Any]) -> Dict[str, nn.Parameter]: + params = { + '_0': init_rand(x, qtz_func=quantization, deqtz_func=dequantization, param='_0', params_list=['_0', '_s'], **kwargs), + } + params['_s'] = init_linear_scale(x, params=params, qtz_func=quantization, **kwargs) + params = {k: nn.Parameter(v, requires_grad=False) for k, v in params.items()} + + if 'post_init_hook' in kwargs: + kwargs['post_init_hook'](parameters=params) + + params = learn_parameters(x, params, + qtz_func=quantization, + deqtz_func=dequantization, + bits=kwargs['bits'], + target_dtype=torch.int8, + epochs=500, + early_stop=False, + ) + if 'post_train_hook' in kwargs: + kwargs['post_train_hook'](parameters=params) + + return params + + +############### Numpy Qtz ############### + + +def np_quantization(x, _0, _s): + return (np.divide(1, np_replace_num(_s, num=0, to=10000)) * np.sin((_0 * x))) + + +def np_dequantization(x, _0, _s): + return (np.divide(1, np_replace_num(_0, num=0, to=10000)) * np.arcsin(np_domain_guard((_s * x), min=-0.99999, max=0.99999, nan=0))) + + +def fit_func(x, _0, _s): + x_ = np_quantization(x, _0, _s) + x_ = np_dequantization(x_, _0, _s) + return x_ + + + +############### HELPERS ############### + +def domain_guard( + x: torch.Tensor, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> torch.Tensor: + """Guard a tensor to a valid domain.""" + x = torch.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = torch.clamp(x, min=min, max=max) + return x + + +def replace_num(x: torch.Tensor, num: float, to: float) -> torch.Tensor: + """Replace a number in a tensor with another number. + + Args: + x (torch.Tensor): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + torch.Tensor: The tensor with the number replaced. + """ + return torch.where(x == num, to, x) + + +def guarded_torch_power(x: torch.Tensor, exp: float) -> torch.Tensor: + """Guard the power operation to a valid domain.""" + return torch.pow(x, exp) if exp >= 1 else torch.pow(torch.relu(x), exp) + + +def init_ones(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.ones_like(val, dtype=torch.float32, device=x.device) + + +def init_rand(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.randn_like(val, dtype=torch.float32, device=x.device) + + +def init_space_search( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + def _build_initial_param(tensor: torch.Tensor, max_initial: int, n_params: int): + """Generates the initial set of parameters. The first iteration generates 10 times more parameters.""" + for _ in range(n_params * 10): # The first iteration generates 10 times more parameters + yield init_rand(tensor) * max_initial # Generates n_params in range [-max_initial, max_initial] + + def _search_param(tensors: List[torch.tensor], n_params): + """Takes the best parameters and generates new parameters around the mean of the best parameters.""" + torch_tensors = torch.stack(tensors) + min_vals, max_vals = torch.aminmax(torch_tensors, dim=0) + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + mean = torch.mean(torch_tensors, dim=0) + for _ in range(n_params): # Generates n_params around the mean of the tensors + yield torch.randn_like(min_vals) * abs_max_val_per_ch + mean + + def _calc(x, qtz_func, deqtz_func, **params): + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params) + x_ = deqtz_func(x=x_, **params) + x_ = x_.transpose(0, 1) + return x_ + + assert "qtz_func" in kwargs, "qtz_func must be provided." + assert "deqtz_func" in kwargs, "deqtz_func must be provided." + assert "params_list" in kwargs, "params list must be provided." + assert "param" in kwargs, "param must be provided." + + qtz_func = kwargs.get('qtz_func') + deqtz_func = kwargs.get('deqtz_func') + params_list = kwargs.get('params_list') + param = kwargs.get('param') + + n_runs = 50 # Number of runs to try to find the best parameters + n_random_params = 50 # Number of random parameters to generate + n_best_to_pick = 5 # Number of best parameters to pick after each run + max_initial = 10000 # Maximum value to initialize the parameters + + # Initializes the parameters + base_params = { p: init_ones(x, **kwargs) for p in params_list if p != param } + params = _build_initial_param(x, max_initial, n_random_params) + + # Performs the search + for _ in range(n_runs): + + best_params = [] + for param_ in params: + try: + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: param_}) + loss_ones = nn.MSELoss()(x, x_) + + if len(best_params) < n_best_to_pick: + best_params.append((param_, loss_ones.item())) + best_params = sorted(best_params, key=lambda x: x[1]) + elif loss_ones < best_params[-1][1]: + best_params[-1] = (param_, loss_ones.item()) + best_params = sorted(best_params, key=lambda x: x[1]) + + except Exception: # The parameters might not be valid for the function's domain + continue + + # Generates new parameters around the mean + params = _search_param([p for p, _ in best_params], n_random_params) + + # Checks if the best parameter is better than the init_ones + p_ones = init_ones(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_ones}) + loss_ones = nn.MSELoss()(x, x_) + + # Checks if the best parameter is better than the init_rand + p_rand = init_rand(x, **kwargs) + x_ = _calc(x, qtz_func, deqtz_func, **base_params, **{param: p_rand}) + loss_rand = nn.MSELoss()(x, x_) + + if loss_rand < best_params[0][1] and loss_rand < loss_ones: + return p_rand + elif loss_ones < best_params[0][1] and loss_ones < loss_rand: + return p_ones + else: + return best_params[0][0] + + +def init_linear_scale( # Symmetric scale. From the study folder + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + assert "bits" in kwargs, "bits must be provided." + assert "params" in kwargs, "params must be provided." + assert "qtz_func" in kwargs, "qtz_func must be provided." + + bits = kwargs.get('bits') + params = kwargs.get('params') + qtz_func = kwargs.get('qtz_func') + + x_ = x.transpose(0, 1) + x_ = qtz_func(x=x_, **params, _s=init_ones(x, **kwargs)) + x_ = x_.transpose(0, 1) + + quant_min, quant_max = get_min_max_from_bits_signed(bits) + min_vals, max_vals = torch.aminmax(x_, dim=1) + min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) + max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) + + eps = torch.finfo(torch.float32).eps + + abs_max_val_per_ch = torch.max(-min_vals, max_vals) + scale = abs_max_val_per_ch / (float(quant_max - quant_min) / 2) + + scale = torch.clamp(scale, min=eps).to(dtype=torch.float32, device=min_vals.device) + + # Introduces some noise in scale + # If I don't introduce noise, the accuracy is going to be 0.0 and not learn anything + # scale = scale + 0.01 * torch.randn_like(scale) + return scale + + +def init_non_linear_regression_fit( + x: torch.Tensor, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + + assert "params_list" in kwargs, "params list must be provided." + assert "np_fit_func" in kwargs, "np_fit_func must be provided." + assert "p0" in kwargs, "p0 must be provided." + np_fit_func = kwargs.get('np_fit_func') + params_list = kwargs.get('params_list') + p0 = kwargs.get('p0') + + def _fit(xdata: np.ndarray, ydata: np.ndarray, func: Callable, p0: List[float]): + popt, _ = curve_fit( + func, + xdata, + ydata, + maxfev=1000, + p0=p0, + method='lm' + ) + return popt + + # 1. Needs to convert the torch tensor to numpy tensor + xdata = x.cpu().numpy() + + # 2. Sorts the data so that it makes it easier to fit to it + sorted_xdata = np.sort(xdata, axis=-1) + + p0 = {k: v.cpu().numpy() for k, v in p0.items()} + params_list = sorted(params_list) # We need to make sure that it matches the numpy fit func arg order + + # 3. Finds the best parameters for each channel + try: + params = [] + for i in range(sorted_xdata.shape[0]): + xdata_ = sorted_xdata[i] + p0_ = [p0[p][i] for p in params_list] + ch_params = _fit(xdata_, xdata_, np_fit_func, p0_) + params.append(ch_params) + + # 4. Builds the parameters + result = {} + for i, p in enumerate(params_list): + result[p] = torch.tensor([p_[i] for p_ in params], dtype=torch.float32).to(x.device) + + return result + + except ValueError as e: + print(f"Could not fit the function with error: {e}") + print(f"Using fallback result...") + return { + k: torch.tensor(v, dtype=torch.float32).to(x.device) for k, v in p0.items() + } + + +def init_zeros(x: torch.Tensor, **kwargs: Dict[str, Any]) -> torch.Tensor: + val = torch.amin(x, dim=1) + return torch.zeros_like(val, dtype=torch.float32, device=x.device) + + +def init_inner_scale(tensor: torch.Tensor, _min: float = torch.inf, _max: float = torch.inf) -> torch.Tensor: + # Calculate the original minimum and maximum values + min_vals, max_vals = torch.aminmax(tensor, dim=-1) + x_min = torch.min(min_vals, torch.zeros_like(min_vals)) + x_max = torch.max(max_vals, torch.zeros_like(max_vals)) + + if _max is torch.inf: # We do not need to scale the tensor. Just need to move it + return torch.ones_like(x_min) + + # Calculate the scale factor + scale = (_max - _min) / (x_max - x_min) + return scale + + + +############## Quant ############### + +@torch.enable_grad() +def learn_parameters( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + qtz_func: nn.Module, + deqtz_func: nn.Module, + bits: int, + target_dtype: torch.dtype, + epochs: int = 1000, + early_stop: bool = True, + do_report: bool = False +) -> Tuple[Dict[str, nn.Parameter], torch.Tensor]: + loss_fn = nn.MSELoss() + + # Determines the initial learning rate by computing the initial loss and multiplying it by + # the order of magnitude of the loss divided by 2 + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + base_lr = 0.1 + exponent = int(np.floor(np.log10(loss.item()))) + lr = base_lr * (10 ** (exponent // 2)) + + # Requires gradients in the parameters + for p in params.values(): + p.requires_grad = True + p.grad = None + + param_keys = list(params.keys()) + param_values = list(params.values()) + + # Defines optimizer and loss function + optimizer = torch.optim.Adam(param_values, lr=lr) + scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs // 10) + + # Contains the best loss and the best parameters + best_loss = float("inf") + best_params = None + + # Used to stop the search early + min_delta = 1e-7 + acc_loss = [] + percent_epochs_before_stop = 0.1 + + for i in range(epochs): + optimizer.zero_grad() + + quant = quantize(x, params, qtz_func, bits, target_dtype) + dequant = dequantize(quant, params, deqtz_func, bits, x.dtype) + loss = loss_fn(x, dequant) + + if loss.isnan() or loss.isinf(): + raise Exception("Loss is NaN or Inf. Stopping the search.") + + loss.backward() + optimizer.step() + scheduler.step() + + acc_loss.append(loss.item()) + + # Reports loss every 10 steps + if i % 10 == 0 and do_report: + print(f"Epoch {i}: Loss {loss.item()}") + + # Optimizes the parameter search by storing the best loss and the parameters + if loss.item() < best_loss: + best_loss = loss.item() + best_params = copy.deepcopy({ + k: v for k, v in params.items() if k in param_keys + }) + + # We also stop the search if the loss has not considerably during the last 10% epochs + if early_stop: + epochs_before_stop = int(epochs * percent_epochs_before_stop) + if i > epochs_before_stop and abs(acc_loss[i - epochs_before_stop] - acc_loss[i]) < min_delta: + break + + # No longer requires gradients in the parameters + for p in best_params.values(): + p.requires_grad = False + p.grad = None + + if do_report: + return best_params, acc_loss + else: + return best_params + + +def quantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + target_dtype: torch.dtype = torch.int8 +) -> torch.Tensor: + quant_min, quant_max = get_min_max_from_bits_signed(bits) + x = x.transpose(0, 1) # Aligns shapes + x = func(x=x, **params) + x = x.transpose(0, 1) + x = torch.clamp(round_func_BPDA(x), quant_min, quant_max).to(target_dtype) + return x + + +def dequantize( + x: torch.Tensor, + params: Dict[str, nn.Parameter], + func: nn.Module, + bits: int, + out_dtype: torch.dtype +) -> torch.Tensor: + x = x.to(dtype=out_dtype) + x = x.transpose(0, 1) + x = func(x=x, **params) + x = x.transpose(0, 1) + return x + + +def round_func_BPDA(input): + # This is equivalent to replacing round function (non-differentiable) with + # an identity function (differentiable) only when backward. + forward_value = torch.round(input) + out = input.clone() + out.data = forward_value.data + return out + + +def get_min_max_from_bits_signed(bit_width: int) -> Tuple[int, int]: + return -2 ** (bit_width - 1), 2 ** (bit_width - 1) - 1 + + + +############## Numpy ############### + +def np_domain_guard( + x: np.ndarray, + min: float = None, + max: float = None, + posinf: float = None, + neginf: float = None, + nan: float = None + ) -> np.ndarray: + """Guard a tensor to a valid domain.""" + x = np.nan_to_num(x, posinf=posinf, neginf=neginf, nan=nan) + if min is not None or max is not None: + x = np.clip(x, min, max) + return x + + +def np_replace_num(x: np.ndarray, num: float, to: float) -> np.ndarray: + """Replace a number in a tensor with another number. + + Args: + x (np.ndarray): The input tensor. + num (float): The number to replace. + to (float): The number to replace with. + + Returns: + np.ndarray: The tensor with the number replaced. + """ + return np.where(x == num, to, x) + + +def np_guarded_power(x: np.ndarray, exp: float) -> np.ndarray: + """Guard the power operation to a valid domain.""" + return np.power(x, exp) if exp >= 1 else np.power(np.maximum(x, 0), exp) + diff --git a/fn_gen/rnd_naive_t/9/loss.png b/fn_gen/rnd_naive_t/9/loss.png new file mode 100644 index 0000000000000000000000000000000000000000..2f980524a41387a9989c73e7007f9c7105844929 Binary files /dev/null and b/fn_gen/rnd_naive_t/9/loss.png differ diff --git a/fn_gen/rnd_naive_t/9/quantization.png b/fn_gen/rnd_naive_t/9/quantization.png new file mode 100644 index 0000000000000000000000000000000000000000..6c19a77ca575b42ec9fd3c837196ba98d11fe711 Binary files /dev/null and b/fn_gen/rnd_naive_t/9/quantization.png differ