concrete-ml-encrypted-deeplearning / creating_models.py

chore: adding a test with DL

6e245b1 5 months ago

No virus

5.9 kB

	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.")