Skip to content

Commit

Permalink
chore: only create keys and initialise executor once
Browse files Browse the repository at this point in the history
  • Loading branch information
jfrery committed Nov 18, 2024
1 parent f0cd673 commit 5f73bd6
Showing 1 changed file with 6 additions and 10 deletions.
16 changes: 6 additions & 10 deletions src/concrete/ml/torch/hybrid_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -456,7 +456,6 @@ def forward(self, x: torch.Tensor, fhe: str = "disable") -> torch.Tensor:

# Validate the FHE mode
fhe_mode = HybridFHEMode(fhe)
self.executor = None

if _HAS_GLWE_BACKEND and self._has_large_linear_layers:
if fhe_mode == HybridFHEMode.SIMULATE:
Expand All @@ -466,15 +465,12 @@ def forward(self, x: torch.Tensor, fhe: str = "disable") -> torch.Tensor:
)

if fhe_mode in (HybridFHEMode.EXECUTE, HybridFHEMode.REMOTE, HybridFHEMode.DISABLE):
# If all layers are pure linear, enable the GLWE optimization for all layers
# and generate an encryption and compression key for all layers
# as they share crypto-parameters

# Loading keys from a file could be done here, and the
# keys could be passed as arguments to the Executor
self.executor = GLWELinearLayerExecutor()
if fhe_mode != HybridFHEMode.DISABLE:
self.executor.keygen()
# Initialize executor only if not already done
if self.executor is None:
self.executor = GLWELinearLayerExecutor()
# Generate keys only if needed and not already done
if fhe_mode != HybridFHEMode.DISABLE:
self.executor.keygen()

# Update executor for all remote modules
for module in self.remote_modules.values():
Expand Down

0 comments on commit 5f73bd6

Please sign in to comment.