mirror of
https://github.com/luxfhe/ml-sdk.git
synced 2026-10-30 18:13:04 +00:00
zeekay
bde2e2531b
Machine Learning with FHE - fork of Zama Concrete-ML - Built-in FHE-friendly ML models - scikit-learn compatible API - PyTorch model conversion - Quantization-aware training
47 lines
19 KiB
JSON
47 lines
19 KiB
JSON
{"text": "class TinyCNN(nn.Module):"}
|
|
{"text": "def __init__(self, n_classes) -> None:\n super().__init__()\n self.conv1 = nn.Conv2d(1, 8, 3, stride=1, padding=0)\n self.conv2 = nn.Conv2d(8, 16, 3, stride=2, padding=0)\n self.conv3 = nn.Conv2d(16, 32, 2, stride=1, padding=0)\n self.fc1 = nn.Linear(32, n_classes)"}
|
|
{"text": "def forward(self, x):\n x = self.conv1(x)\n x = torch.relu(x)\n x = self.conv2(x)\n x = torch.relu(x)\n x = self.conv3(x)\n x = torch.relu(x)\n x = x.flatten(1)\n x = self.fc1(x)\n return x\nnet = TinyCNN(10)\n#... (Training loop)...\nq_module = compile_torch_model(net, x_train, rounding_threshold_bits=6, p_error=0.1)\n# Key generation\nq_module.fhe_circuit.keygen()\n# Inference"}
|
|
{"text": "self.fc1(x)\n return x\nnet = TinyCNN(10)\n#... (Training loop)...\nq_module = compile_torch_model(net, x_train, rounding_threshold_bits=6, p_error=0.1)\n# Key generation\nq_module.fhe_circuit.keygen()\n# Inference in FHE\ny_pred_fhe = q_module.forward(x_test, fhe=\"execute\")\n\n**4. Quantization-Aware Training:**\npython\nfrom torch import nn\nfrom concrete.ml.torch.compile import compile_brevitas_qat_model\nimport brevitas.nn as qnn\nfrom brev"}
|
|
{"text": "in FHE\ny_pred_fhe = q_module.forward(x_test, fhe=\"execute\")\n\n**4. Quantization-Aware Training:**\npython\nfrom torch import nn\nfrom concrete.ml.torch.compile import compile_brevitas_qat_model\nimport brevitas.nn as qnn\nfrom brevitas.core.bit_width import BitWidthImplType\nfrom brevitas.core.quant import QuantType\nfrom brevitas.core.restrict_val import FloatToIntImplType, RestrictValueType\nfrom brevitas.core.scaling import ScalingImplType\nfrom brevitas.core.zero_point import ZeroZeroPoint\nfrom bre"}
|
|
{"text": "itas.core.bit_width import BitWidthImplType\nfrom brevitas.core.quant import QuantType\nfrom brevitas.core.restrict_val import FloatToIntImplType, RestrictValueType\nfrom brevitas.core.scaling import ScalingImplType\nfrom brevitas.core.zero_point import ZeroZeroPoint\nfrom brevitas.inject import ExtendedInjector\nfrom brevitas.quant.solver import ActQuantSolver, WeightQuantSolver\nfrom dependencies import value\nfrom torch.nn.utils import prune\n#... (Data loading and preprocessing)..."}
|
|
{"text": "class CommonQuant(ExtendedInjector):\n bit_width_impl_type = BitWidthImplType.CONST\n scaling_impl_type = ScalingImplType.CONST\n restrict_scaling_type = RestrictValueType.FP\n zero_point_impl = ZeroZeroPoint\n float_to_int_impl_type = FloatToIntImplType.ROUND\n scaling_per_output_channel = False\n narrow_range = True\n signed = True\n @value"}
|
|
{"text": "def quant_type(bit_width): # pylint: disable=no-self-argument\n if bit_width is None:\n return QuantType.FP\n if bit_width == 1:\n return QuantType.BINARY\n return QuantType.INT"}
|
|
{"text": "class CommonWeightQuant(CommonQuant, WeightQuantSolver): # pylint: disable=too-many-ancestors\n scaling_const = 1.0\n signed = True"}
|
|
{"text": "class CommonActQuant(CommonQuant, ActQuantSolver): # pylint: disable=too-many-ancestors\n min_val = -1.0\n max_val = 1.0"}
|
|
{"text": "class QATPrunedSimpleNet(nn.Module):"}
|
|
{"text": "def __init__(self, n_hidden, qlinear_args, qidentity_args):\n super().__init__()\n self.pruned_layers = set()\n self.quant_inp = qnn.QuantIdentity(**qidentity_args)\n self.fc1 = qnn.QuantLinear(IN_FEAT, n_hidden, **qlinear_args)\n self.relu1 = qnn.QuantReLU(bit_width=qidentity_args[\"bit_width\"])\n self.fc2 = qnn.QuantLinear(n_hidden, n_hidden, **qlinear_args)\n self.relu2 = qnn.QuantReLU(bit_width=qidentity_args[\"bit_width"}
|
|
{"text": ", **qlinear_args)\n self.relu1 = qnn.QuantReLU(bit_width=qidentity_args[\"bit_width\"])\n self.fc2 = qnn.QuantLinear(n_hidden, n_hidden, **qlinear_args)\n self.relu2 = qnn.QuantReLU(bit_width=qidentity_args[\"bit_width\"])\n self.fc3 = qnn.QuantLinear(n_hidden, OUT_FEAT, **qlinear_args)\n for m in self.modules():\n if isinstance(m, qnn.QuantLinear):\n torch.nn.init.uniform_(m.weight.data, -1, 1)"}
|
|
{"text": "def forward(self, x):\n x = self.quant_inp(x)\n x = self.relu1(self.fc1(x))\n x = self.relu2(self.fc2(x))\n x = self.fc3(x)\n return x"}
|
|
{"text": "def prune(self, max_non_zero):\n # Linear layer weight has dimensions NumOutputs x NumInputs\n for name, layer in self.named_modules():\n if isinstance(layer, qnn.QuantLinear):\n num_zero_weights = (layer.weight.shape[1] - max_non_zero) * layer.weight.shape[0]\n if num_zero_weights <= 0:\n continue\n print(f\"Pruning layer {name} factor {num_zero_weights}\")\n prune.l1_unstructured(layer, \"weight\", amount=num_zero_weights)\n self.pruned_layers.add(name)"}
|
|
{"text": "def unprune(self):\n for name, layer in self.named_modules():\n if name in self.pruned_layers:\n prune.remove(layer, \"weight\")\n self.pruned_layers.remove(name)\ntorch_model = QATPrunedSimpleNet(\n n_hidden=n_hidden,\n qlinear_args={\n \"weight_bit_width\": 3,\n \"weight_quant\": CommonWeightQuant,\n \"bias\": True,\n \"bias_quant\": None,\n \"narrow_range\": True,\n },\n qidentity_args={\"bit_width\": 3, \"act_quant\": CommonActQuant},\n)\ntorch"}
|
|
{"text": "_args={\n \"weight_bit_width\": 3,\n \"weight_quant\": CommonWeightQuant,\n \"bias\": True,\n \"bias_quant\": None,\n \"narrow_range\": True,\n },\n qidentity_args={\"bit_width\": 3, \"act_quant\": CommonActQuant},\n)\ntorch_model.prune(20)\n#... (Training loop)...\nquantized_numpy_module = compile_brevitas_qat_model(torch_model, x_train)\n# Inference in FHE (simulation)\ny_pred_fhe = quantized_numpy_module.forward(x_test, fhe=\"simulate\")\n\n**5. Client/Server"}
|
|
{"text": "_model.prune(20)\n#... (Training loop)...\nquantized_numpy_module = compile_brevitas_qat_model(torch_model, x_train)\n# Inference in FHE (simulation)\ny_pred_fhe = quantized_numpy_module.forward(x_test, fhe=\"simulate\")\n\n**5. Client/Server Deployment (LogisticRegressionTraining.ipynb):**\npython\nfrom pathlib import Path\nfrom tempfile import TemporaryDirectory\nimport numpy as np\nfrom concrete.ml.deployment import FHEModelClient, FHEModelDev, FHEModelServer\nfrom concrete.ml.sklearn import SGDClassifier\nfrom concrete import fhe"}
|
|
{"text": "Deployment (LogisticRegressionTraining.ipynb):**\npython\nfrom pathlib import Path\nfrom tempfile import TemporaryDirectory\nimport numpy as np\nfrom concrete.ml.deployment import FHEModelClient, FHEModelDev, FHEModelServer\nfrom concrete.ml.sklearn import SGDClassifier\nfrom concrete import fhe\n#... (Data loading, preprocessing, and model training)...\n# Assuming you have a trained model: sgd_clf_binary_fhe\n# and x_compile_set, y_compile_set for compilation\n# Define the directory where to save the deployment files\nDEPLOYMENT_PATH = Path(\"fhe_training\")"}
|
|
{"text": "#... (Data loading, preprocessing, and model training)...\n# Assuming you have a trained model: sgd_clf_binary_fhe\n# and x_compile_set, y_compile_set for compilation\n# Define the directory where to save the deployment files\nDEPLOYMENT_PATH = Path(\"fhe_training\")\nDEPLOYMENT_PATH.mkdir(exist_ok=True)\ndeployment_dir = TemporaryDirectory(dir=str(DEPLOYMENT_PATH))\ndeployment_path = Path(deployment_dir.name)\n# Save the model for deployment\nfhe_dev = FHEModelDev(deployment_path, sgd_clf_binary_fhe)\nfhe_dev.save(mode=\""}
|
|
{"text": "DEPLOYMENT_PATH.mkdir(exist_ok=True)\ndeployment_dir = TemporaryDirectory(dir=str(DEPLOYMENT_PATH))\ndeployment_path = Path(deployment_dir.name)\n# Save the model for deployment\nfhe_dev = FHEModelDev(deployment_path, sgd_clf_binary_fhe)\nfhe_dev.save(mode=\"training\")\n# Client-side setup\nfhe_client = FHEModelClient(deployment_path)\nfhe_client.load()\nserialized_evaluation_keys = fhe_client.get_serialized_evaluation_keys()\n# Server-side setup\nfhe_server = FHEModelServer(deployment_path)\nfhe_server.load()\n# Example of encryption,"}
|
|
{"text": "training\")\n# Client-side setup\nfhe_client = FHEModelClient(deployment_path)\nfhe_client.load()\nserialized_evaluation_keys = fhe_client.get_serialized_evaluation_keys()\n# Server-side setup\nfhe_server = FHEModelServer(deployment_path)\nfhe_server.load()\n# Example of encryption, server-side processing, and decryption\nbatch_size = sgd_clf_binary_fhe.batch_size\nweights = np.random.rand(1, x_train.shape[1], 1)\nbias = np.random.rand(1, 1, 1)"}
|
|
{"text": "def quantize_encrypt_serialize_batches(fhe_client, x, y, weights, bias, batch_size):\n #... (Implementation as before)..."}
|
|
{"text": "def server_run(fhe_server, x_batches_enc, y_batches_enc, weights_enc, bias_enc, evaluation_keys):\n #... (Implementation as before)..."}
|
|
{"text": "def train_fhe_client_server(\n #... (Parameters as before)...\n):\n #... (Training loop)\n # Quantize, encrypt and serialize the batched inputs as well as the weight and bias values\n x_batches_enc, y_batches_enc, weights_enc, bias_enc = quantize_encrypt_serialize_batches(\n fhe_client, x, y, weights, bias, batch_size\n )\n # Iterate the circuit over the batches on the server\n fitted_weights_enc, fitted_bias_enc = server_run(\n fhe_server,\n x_batches_enc,\n y_batches_enc,\n weights_enc,"}
|
|
{"text": "_serialize_batches(\n fhe_client, x, y, weights, bias, batch_size\n )\n # Iterate the circuit over the batches on the server\n fitted_weights_enc, fitted_bias_enc = server_run(\n fhe_server,\n x_batches_enc,\n y_batches_enc,\n weights_enc,\n bias_enc,\n serialized_evaluation_keys,\n )\n # Back on the client, deserialize, decrypt and de-quantize the fitted weight and bias values\n weights, bias = fhe_client.deserialize_decrypt_dequantize(\n fitted_weights_enc, fitted_bias_enc\n )\n return weights, bias,"}
|
|
{"text": "bias_enc,\n serialized_evaluation_keys,\n )\n # Back on the client, deserialize, decrypt and de-quantize the fitted weight and bias values\n weights, bias = fhe_client.deserialize_decrypt_dequantize(\n fitted_weights_enc, fitted_bias_enc\n )\n return weights, bias, acc_history\n# Cleanup\ndeployment_dir.cleanup()\n\n**6. Hyper-parameter Tuning with GridSearchCV (XGBClassifier.ipynb, DecisionTreeRegressor.ipynb):**\npython\nfrom sklearn.model_selection import GridSearchCV\nfrom concrete.ml.sklearn import XGBClassifier as ConcreteXGBClassifier\nfrom"}
|
|
{"text": "acc_history\n# Cleanup\ndeployment_dir.cleanup()\n\n**6. Hyper-parameter Tuning with GridSearchCV (XGBClassifier.ipynb, DecisionTreeRegressor.ipynb):**\npython\nfrom sklearn.model_selection import GridSearchCV\nfrom concrete.ml.sklearn import XGBClassifier as ConcreteXGBClassifier\nfrom sklearn.metrics import make_scorer, matthews_corrcoef\n#... (Data loading and preprocessing)...\n# Create scorer with the MCC metric\ngrid_scorer = make_scorer(matthews_corrcoef, greater_is_better=True)\n# Define the parameter grid to search\nparam_grid = {"}
|
|
{"text": "sklearn.metrics import make_scorer, matthews_corrcoef\n#... (Data loading and preprocessing)...\n# Create scorer with the MCC metric\ngrid_scorer = make_scorer(matthews_corrcoef, greater_is_better=True)\n# Define the parameter grid to search\nparam_grid = {\n \"n_bits\": [5, 6],\n \"max_depth\": [2, 3],\n \"n_estimators\": [10, 20, 50],\n}\n# Instantiate GridSearchCV with the Concrete ML model\ngrid_search = GridSearchCV(\n ConcreteXGBClassifier(),\n param_grid"}
|
|
{"text": "\"n_bits\": [5, 6],\n \"max_depth\": [2, 3],\n \"n_estimators\": [10, 20, 50],\n}\n# Instantiate GridSearchCV with the Concrete ML model\ngrid_search = GridSearchCV(\n ConcreteXGBClassifier(),\n param_grid,\n cv=5,\n scoring=grid_scorer,\n error_score=\"raise\",\n verbose=1,\n)\n# Run the grid search\ngrid_search.fit(x_train, y_train)\n# Get the best parameters\nbest_params = grid_search.best_params_\n# Create a new model with the best parameters"}
|
|
{"text": ",\n cv=5,\n scoring=grid_scorer,\n error_score=\"raise\",\n verbose=1,\n)\n# Run the grid search\ngrid_search.fit(x_train, y_train)\n# Get the best parameters\nbest_params = grid_search.best_params_\n# Create a new model with the best parameters\nbest_model = ConcreteXGBClassifier(**best_params)\nbest_model.fit(x_train, y_train)\n# Compile and proceed with FHE inference as shown in other examples\n\n**7. GLM Models (GLMComparison.ipynb):**\n* **Poisson Regressor**\npython\nfrom concrete"}
|
|
{"text": "best_model = ConcreteXGBClassifier(**best_params)\nbest_model.fit(x_train, y_train)\n# Compile and proceed with FHE inference as shown in other examples\n\n**7. GLM Models (GLMComparison.ipynb):**\n* **Poisson Regressor**\npython\nfrom concrete.ml.sklearn import PoissonRegressor as ConcretePoissonRegressor\n#... (Data loading and preprocessing)...\nconcrete_pr = ConcretePoissonRegressor(n_bits=8)\nconcrete_pr.fit(x_train, y_train, sample_weight=train_weights)\ncircuit = concrete_pr.compile(x_train)\n# Key generation"}
|
|
{"text": ".ml.sklearn import PoissonRegressor as ConcretePoissonRegressor\n#... (Data loading and preprocessing)...\nconcrete_pr = ConcretePoissonRegressor(n_bits=8)\nconcrete_pr.fit(x_train, y_train, sample_weight=train_weights)\ncircuit = concrete_pr.compile(x_train)\n# Key generation\ncircuit.client.keygen(force=False)\n# Inference in FHE\ny_pred_fhe = concrete_pr.predict(x_test, fhe=\"execute\")\n\n* **Gamma Regressor**\npython\nfrom concrete.ml.sklearn import GammaRegressor as ConcreteGammaRegressor\n#... (Data loading and preprocessing)..."}
|
|
{"text": "circuit.client.keygen(force=False)\n# Inference in FHE\ny_pred_fhe = concrete_pr.predict(x_test, fhe=\"execute\")\n\n* **Gamma Regressor**\npython\nfrom concrete.ml.sklearn import GammaRegressor as ConcreteGammaRegressor\n#... (Data loading and preprocessing)...\nconcrete_gr = ConcreteGammaRegressor(n_bits=8)\nconcrete_gr.fit(x_train, y_train, sample_weight=train_weights)\ncircuit = concrete_gr.compile(x_train)\n# Key generation\ncircuit.client.keygen(force=False)\n# Inference in FHE\ny_pred_fhe = concrete_gr.predict(x"}
|
|
{"text": "concrete_gr = ConcreteGammaRegressor(n_bits=8)\nconcrete_gr.fit(x_train, y_train, sample_weight=train_weights)\ncircuit = concrete_gr.compile(x_train)\n# Key generation\ncircuit.client.keygen(force=False)\n# Inference in FHE\ny_pred_fhe = concrete_gr.predict(x_test, fhe=\"execute\")\n\n* **Tweedie Regressor**\npython\nfrom concrete.ml.sklearn import TweedieRegressor as ConcreteTweedieRegressor\n#... (Data loading and preprocessing)...\nconcrete_tr = ConcreteTweedieRegressor(n_bits=8, power=1.9"}
|
|
{"text": "_test, fhe=\"execute\")\n\n* **Tweedie Regressor**\npython\nfrom concrete.ml.sklearn import TweedieRegressor as ConcreteTweedieRegressor\n#... (Data loading and preprocessing)...\nconcrete_tr = ConcreteTweedieRegressor(n_bits=8, power=1.9)\nconcrete_tr.fit(x_train, y_train, sample_weight=train_weights)\ncircuit = concrete_tr.compile(x_train)\n# Key generation\ncircuit.client.keygen(force=False)\n# Inference in FHE\ny_pred_fhe = concrete_tr.predict(x_test, fhe=\"execute\")\n\n**8. Fine"}
|
|
{"text": ")\nconcrete_tr.fit(x_train, y_train, sample_weight=train_weights)\ncircuit = concrete_tr.compile(x_train)\n# Key generation\ncircuit.client.keygen(force=False)\n# Inference in FHE\ny_pred_fhe = concrete_tr.predict(x_test, fhe=\"execute\")\n\n**8. Fine-tuning with LoRA (LoraMLP.ipynb):**\npython\nimport torch\nfrom peft import LoraConfig, get_peft_model\nfrom torch import nn, optim\nfrom concrete.ml.torch.lora import LoraTrainer\n#... (Data loading and preprocessing)...\n# Define"}
|
|
{"text": "-tuning with LoRA (LoraMLP.ipynb):**\npython\nimport torch\nfrom peft import LoraConfig, get_peft_model\nfrom torch import nn, optim\nfrom concrete.ml.torch.lora import LoraTrainer\n#... (Data loading and preprocessing)...\n# Define an MLP model without LoRA layers"}
|
|
{"text": "class SimpleMLP(nn.Module):"}
|
|
{"text": "def __init__(self, input_size=2, hidden_size=128, num_classes=2):\n super().__init__()\n self.fc1 = nn.Linear(input_size, hidden_size)\n self.relu = nn.ReLU()\n self.fc2 = nn.Linear(hidden_size, num_classes)"}
|
|
{"text": "def forward(self, x):\n out = self.fc1(x)\n out = self.relu(out)\n out = self.fc2(out)\n return out\n# Instantiate the model\nmodel = SimpleMLP()\n#... (Training loop for Task 1)...\n# Apply LoRA to the model using peft\nlora_config = LoraConfig(\n r=1, lora_alpha=1, lora_dropout=0.01, target_modules=[\"fc1\", \"fc2\"], bias=\"none\"\n)\npeft_model = get_peft_model(model, lora_config)\n# Update training parameters"}
|
|
{"text": "using peft\nlora_config = LoraConfig(\n r=1, lora_alpha=1, lora_dropout=0.01, target_modules=[\"fc1\", \"fc2\"], bias=\"none\"\n)\npeft_model = get_peft_model(model, lora_config)\n# Update training parameters, including loss function\noptimizer = optim.Adam(filter(lambda p: p.requires_grad, peft_model.parameters()), lr=0.01)\nloss_fn = nn.CrossEntropyLoss()\ntraining_args = {\"gradient_accumulation_steps\": 1}\n# Set up LoRA training\nlora_trainer = LoraTrainer"}
|
|
{"text": ", including loss function\noptimizer = optim.Adam(filter(lambda p: p.requires_grad, peft_model.parameters()), lr=0.01)\nloss_fn = nn.CrossEntropyLoss()\ntraining_args = {\"gradient_accumulation_steps\": 1}\n# Set up LoRA training\nlora_trainer = LoraTrainer(peft_model, optimizer=optimizer, loss_fn=loss_fn, training_args=training_args)\n# Prepare input data for calibration\nbatch_size_per_task = batch_size // 2\ninputset = (\n torch.cat([X_task1[:batch_size_per_task], X_task2[:batch_size_per_task]]"}
|
|
{"text": "(peft_model, optimizer=optimizer, loss_fn=loss_fn, training_args=training_args)\n# Prepare input data for calibration\nbatch_size_per_task = batch_size // 2\ninputset = (\n torch.cat([X_task1[:batch_size_per_task], X_task2[:batch_size_per_task]]),\n torch.cat([y_task1[:batch_size_per_task], y_task2[:batch_size_per_task]]),\n)\n# Compile the model\nlora_trainer.compile(inputset, n_bits=8)\n# Fine-tune the model on Task 2 using LoRA\nlora_trainer.train(train_loader"}
|
|
{"text": "),\n torch.cat([y_task1[:batch_size_per_task], y_task2[:batch_size_per_task]]),\n)\n# Compile the model\nlora_trainer.compile(inputset, n_bits=8)\n# Fine-tune the model on Task 2 using LoRA\nlora_trainer.train(train_loader_task2, num_epochs=10, fhe=\"execute\")\n# Enable/Disable LoRA adapters\npeft_model.enable_adapter_layers()\npeft_model.disable_adapter_layers()\n# Print trainable (lora) parameters\npeft_model.print_trainable_parameters()\n# Save the model and remove all layers that will be done"}
|
|
{"text": "_task2, num_epochs=10, fhe=\"execute\")\n# Enable/Disable LoRA adapters\npeft_model.enable_adapter_layers()\npeft_model.disable_adapter_layers()\n# Print trainable (lora) parameters\npeft_model.print_trainable_parameters()\n# Save the model and remove all layers that will be done on the server\npath = Path(\"lora_mlp\")\nif path.is_dir() and any(path.iterdir()):\n shutil.rmtree(path)\nlora_trainer.save_and_clear_private_info(path)"}
|