File size: 5,896 Bytes
6e245b1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 |
import shutil
import sys
from pathlib import Path
from concrete.ml.deployment import FHEModelDev
from concrete.ml.deployment import FHEModelClient
def compile_and_make_it_deployable(model_dev, X_train):
path_to_model = Path("compiled_model")
n_bits = 4
model_dev = compile_torch_model(model_dev, X_train, rounding_threshold_bits=6, p_error=0.1)
# Accuracy in simulation
accs = test_with_concrete(
model_dev,
test_dataloader,
use_sim=True,
)
print(f"Simulated FHE execution for {n_bits} bit model_devwork accuracy: {accs:.2f}")
# Saving the model
shutil.rmtree(path_to_model, ignore_errors=True)
fhemodel_dev = FHEModelDev(path_to_model, model_dev)
fhemodel_dev.save(via_mlir=True)
# BEGIN: insert your ML task here
# Typically
import time
import numpy as np
import torch
import torch.utils
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm
from concrete.ml.torch.compile import compile_torch_model
X, y = load_digits(return_X_y=True)
# The sklearn Digits data-set, though it contains digit images, keeps these images in vectors
# so we need to reshape them to 2D first. The images are 8x8 px in size and monochrome
X = np.expand_dims(X.reshape((-1, 8, 8)), 1)
X_train, X_test, Y_train, Y_test = train_test_split(
X, y, test_size=0.25, shuffle=True, random_state=42
)
class TinyCNN(nn.Module):
"""A very small CNN to classify the sklearn digits data-set."""
def __init__(self, n_classes) -> None:
"""Construct the CNN with a configurable number of classes."""
super().__init__()
# This model_devwork has a total complexity of 1216 MAC
self.conv1 = nn.Conv2d(1, 8, 3, stride=1, padding=0)
self.conv2 = nn.Conv2d(8, 16, 3, stride=2, padding=0)
self.conv3 = nn.Conv2d(16, 32, 2, stride=1, padding=0)
self.fc1 = nn.Linear(32, n_classes)
def forward(self, x):
"""Run inference on the tiny CNN, apply the decision layer on the reshaped conv output."""
x = self.conv1(x)
x = torch.relu(x)
x = self.conv2(x)
x = torch.relu(x)
x = self.conv3(x)
x = torch.relu(x)
x = x.flatten(1)
x = self.fc1(x)
return x
torch.manual_seed(42)
def train_one_epoch(model_dev, optimizer, train_loader):
# Cross Entropy loss for classification when not using a softmax layer in the model_devwork
loss = nn.CrossEntropyLoss()
model_dev.train()
avg_loss = 0
for data, target in train_loader:
optimizer.zero_grad()
output = model_dev(data)
loss_model_dev = loss(output, target.long())
loss_model_dev.backward()
optimizer.step()
avg_loss += loss_model_dev.item()
return avg_loss / len(train_loader)
def test_torch(model_dev, test_loader):
"""Test the model_devwork: measure accuracy on the test set."""
# Freeze normalization layers
model_dev.eval()
all_y_pred = np.zeros((len(test_loader)), dtype=np.int64)
all_targets = np.zeros((len(test_loader)), dtype=np.int64)
# Iterate over the batches
idx = 0
for data, target in test_loader:
# Accumulate the ground truth labels
endidx = idx + target.shape[0]
all_targets[idx:endidx] = target.numpy()
# Run forward and get the predicted class id
output = model_dev(data).argmax(1).detach().numpy()
all_y_pred[idx:endidx] = output
idx += target.shape[0]
# Print out the accuracy as a percentage
n_correct = np.sum(all_targets == all_y_pred)
print(
f"Test accuracy for fp32 weights and activations: "
f"{n_correct / len(test_loader) * 100:.2f}%"
)
def test_with_concrete(quantized_module, test_loader, use_sim):
"""Test a neural model_devwork that is quantized and compiled with Concrete ML."""
# Casting the inputs into int64 is recommended
all_y_pred = np.zeros((len(test_loader)), dtype=np.int64)
all_targets = np.zeros((len(test_loader)), dtype=np.int64)
# Iterate over the test batches and accumulate predictions and ground truth labels in a vector
idx = 0
for data, target in tqdm(test_loader):
data = data.numpy()
target = target.numpy()
fhe_mode = "simulate" if use_sim else "execute"
# Quantize the inputs and cast to appropriate data type
y_pred = quantized_module.forward(data, fhe=fhe_mode)
endidx = idx + target.shape[0]
# Accumulate the ground truth labels
all_targets[idx:endidx] = target
# Get the predicted class id and accumulate the predictions
y_pred = np.argmax(y_pred, axis=1)
all_y_pred[idx:endidx] = y_pred
# Update the index
idx += target.shape[0]
# Compute and report results
n_correct = np.sum(all_targets == all_y_pred)
return n_correct / len(test_loader)
# Create the tiny CNN with 10 output classes
N_EPOCHS = 50
# Create a train data loader
train_dataset = TensorDataset(torch.Tensor(X_train), torch.Tensor(Y_train))
train_dataloader = DataLoader(train_dataset, batch_size=64)
# Create a test data loader to supply batches for model_devwork evaluation (test)
test_dataset = TensorDataset(torch.Tensor(X_test), torch.Tensor(Y_test))
test_dataloader = DataLoader(test_dataset)
# Train the model_devwork with Adam, output the test set accuracy every epoch
model_dev = TinyCNN(10)
losses_bits = []
optimizer = torch.optim.Adam(model_dev.parameters())
for _ in tqdm(range(N_EPOCHS), desc="Training"):
losses_bits.append(train_one_epoch(model_dev, optimizer, train_dataloader))
# END: insert your ML task here
compile_and_make_it_deployable(model_dev, X_train)
print("Your model is ready to be deployable.")
|