Environment setup

Configure PyTorch precision and suppress non-critical warnings before any model imports. float32_matmul_precision='high' enables TF32 on Ampere-class GPUs, giving a significant training speedup with negligible accuracy loss.

[1]:

import torch
torch.set_float32_matmul_precision('high')
import warnings
warnings.filterwarnings("ignore")

8edaa38550ea404b8e530ddf281997e2

Training Your Custom HOOPS Embeddings Model

The EmbeddingFlowModel is a specialized FlowModel implementation designed for training shape-embedding models from CAD data using self-supervised learning.

This enables data scientists to train custom HOOPS embeddings models on their own CAD datasets.

This notebook is Part 2 of 2 in the custom HOOPS Embeddings workflow:

  1. `demo_HOOPS_EMBEDDINGS_DataPrep.ipynb <./demo_HOOPS_EMBEDDINGS_DataPrep.ipynb>`__ — gather and encode CAD files into a reusable .dataset / .infoset (run once per dataset).

  2. ``demo_HOOPS_EMBEDDINGS_training.ipynb`` (this notebook) — split the encoded dataset, then train and save your model.

Data preparation is separated because it only needs to run once, whereas training is typically repeated many times while tuning hyperparameters and iterating toward convergence. Run the data prep notebook before this one — the steps below load the artifacts it produces.

[2]:

import hoops_ai
import os
import sys

license_key = os.environ.get("HOOPS_AI_LICENSE")
if not license_key:
    sys.exit("HOOPS_AI_LICENSE environment variable is required.")

hoops_ai.set_license(license_key, validate=True)

------------------------------------------------------------
HOOPS AI
------------------------------------------------------------
  Platform      : Windows 11
  Architecture  : AMD64
  Python        : 3.12.10
------------------------------------------------------------
  Core          : hoops-ai             1.1.0  (build: ed23c844 2026-06-12T14:14:13Z)
  CAD Access    : hoops-exchange       26.2.0  (build: 1e11169 2026-06-12T10:38:16Z)
  Conversion    : hoops-converter      26.1.1  (build: 00dc9f6 2026-06-12T10:22:46Z)
  Insights      : hoops-web-viewer     26.1.1  (build: d30058f 2026-06-12T10:22:25Z)
------------------------------------------------------------
======================================================================
[OK] HOOPS AI License: Valid
======================================================================

Prerequisite: encoded dataset

This notebook assumes you have already run `demo_HOOPS_EMBEDDINGS_DataPrep.ipynb <./demo_HOOPS_EMBEDDINGS_DataPrep.ipynb>`__, which encodes your CAD files and writes a .dataset store plus an .infoset metadata file under the flow output directory.

The cell below imports flows_outputdir and get_flow_name from scripts/cad_tasks_embeddings.py — the same module the data prep notebook used — so the paths resolve to exactly the artifacts produced there. As long as the data prep notebook has been run, you can re-run this training notebook repeatedly without redoing the encoding step.

Step 2: Split the dataset

Split the encoded dataset into training, validation, and test sets. The split is stratified on face_types (the distribution of B-rep face categories) so each partition reflects the full variety of the dataset. Adjust the ratios to match the size of your dataset — larger datasets can afford smaller validation and test slices.

[3]:

import pathlib
from scripts.cad_tasks_embeddings import flows_outputdir, get_flow_name

flow_name = get_flow_name()

flow_root_dir = flows_outputdir.joinpath("flows", flow_name)
myFlow_info        = str(flow_root_dir.joinpath(f"{flow_name}.infoset"))
myFlow_dataset     = str(flow_root_dir.joinpath(f"{flow_name}.dataset"))

[4]:

from hoops_ai.dataset import DatasetLoader

# Load the already encoded dataset and perform the split
datasetLoader = DatasetLoader(merged_store_path = myFlow_dataset, parquet_file_path=myFlow_info)

datasetLoader.split(key='random', group="faces",train=0.33, validation=0.33, test=0.34)

INFO:To route to workers diagnostics web server please install jupyter-server-proxy: python -m pip install jupyter-server-proxy
INFO:State start
INFO:  Scheduler at:     tcp://127.0.0.1:63857
INFO:  dashboard at:  http://127.0.0.1:8787/status
INFO:Registering Worker plugin shuffle
INFO:        Start Nanny at: 'tcp://127.0.0.1:63860'
INFO:Register worker <WorkerState 'tcp://127.0.0.1:63864', name: 0, status: init, memory: 0, processing: 0>
INFO:Starting worker compute stream, tcp://127.0.0.1:63864
INFO:Starting established connection to tcp://127.0.0.1:63866
INFO:Receive client connection: Client-f387ed13-80e5-11f1-bba0-f4289de57fc2
INFO:Starting established connection to tcp://127.0.0.1:63867

[DatasetExplorer] Default local cluster started: <Client: 'tcp://127.0.0.1:63857' processes=1 threads=16, memory=7.45 GiB>

DEBUG: Successfully built file lists with 35808 files out of 35808 original file codes

============================================================
DATASET STRUCTURE OVERVIEW
============================================================

Group: edges
------------------------------
  edge_convexities: (4305385,) (int32)
  edge_dihedral_angles: (4305385,) (float32)
  edge_indices: (4305385,) (int32)
  edge_lengths: (4305385,) (float32)
  edge_types: (4305385,) (int32)
  edge_u_grids: (4305385, 10, 6) (float32)
  file_id_code_edges: (4305385,) (int64)

Group: faces
------------------------------
  face_areas: (1700552,) (float32)
  face_centroids: (1700552, 3) (float32)
  face_discretization: (1700552, 100, 7) (float32)
  face_indices: (1700552,) (int32)
  face_loops: (1700552,) (int32)
  face_types: (1700552,) (int32)
  file_id_code_faces: (1700552,) (int64)

Group: graph
------------------------------
  edges_destination: (4305385,) (int32)
  edges_source: (4305385,) (int32)
  file_id_code_graph: (4305385,) (int64)
  num_nodes: (4305385,) (int32)

============================================================

[4]:

(11817, 11816, 12175)

Step 3: Train and save the model (SIGNAL architecture)

EmbeddingFlowModel and FlowTrainer are imported from hoops_ai.ml.EXPERIMENTAL. The EXPERIMENTAL namespace signals that the training API may evolve between releases; the inference API (HOOPSEmbeddings, CADSearch) is stable.

FlowTrainer wraps PyTorch Lightning and handles the training loop, checkpointing, and logging. The checkpoint written to trained_model_path at the end of training can be registered directly with HOOPSEmbeddings.register_model() and used for similarity search exactly like the pre-trained model bundled in the release.

[ ]:

from hoops_ai.ml.EXPERIMENTAL import EmbeddingFlowModel, FlowTrainer

# Pretrained 2M-part SIGNAL-preview checkpoint to warm-start (fine-tune) from.
resume_checkpoint = pathlib.Path.cwd().parent.joinpath(
    "packages", "trained_ml_models", "ts3d_2M_hoops_embeddings_SIGNAL-preview.ckpt")


EmbeddingModel = EmbeddingFlowModel(
    result_dir = str(flow_root_dir),
    log_file   = str(flow_root_dir.joinpath("flow.log")),
    load_checkpoint_using_nn_module = True,
    temp_init = 0.07,
    temp_min = 0.1,
    temp_max = 0.15,
)

# FlowTrainer
flow_trainer = FlowTrainer(
    flowmodel       = EmbeddingModel,
    datasetLoader   = datasetLoader,
    experiment_name = "HOOPS_AI_train_SIGNAL_embeddings",
    result_dir      = flow_root_dir,
    accelerator     = 'gpu',
    devices         = [0], #[0]
    max_epochs      = 10, # increasing this value takes longer but better results
    batch_size      = 8,   # parts per training step
    resume_checkpoint_path = str(resume_checkpoint),
)

trained_model_path = flow_trainer.train()

INFO:GPU available: True (cuda), used: True
INFO:TPU available: False, using: 0 TPU cores

HOOPS Embedding Model
-----------------------------------------------------------------------------------
HOOPS Embedding Model - TRAINING STEP
-----------------------------------------------------------------------------------
Training batch size               : 8
Learning rate                     : (from checkpoint)
Train shuffle                     : True
Train seed                        : None

Train set contains                : 11817 samples (33.00%)
Validation set contains           : 11816 samples (33.00%)
Test set contains                 : 12175 samples (34.00%)
Total samples                     : 35808
Max Epoch                         : 10
Resume checkpoint                 : c:\Users\LuisSalazar.LY-LS-LEGION\Documents\repos\HOOPS-AI-tutorials\packages\trained_ml_models\ts3d_2M_hoops_embeddings_SIGNAL-preview.ckpt

The trained model: c:\Users\LuisSalazar.LY-LS-LEGION\Documents\repos\HOOPS-AI-tutorials\embeddings_pipeline\out\flows\HOOPS_Embedding_Training\flowtrainer\HOOPS_AI_train_SIGNAL_embeddings\0716\091416\best.ckpt

To monitor the logs, run:
tensorboard --logdir c:\Users\LuisSalazar.LY-LS-LEGION\Documents\repos\HOOPS-AI-tutorials\embeddings_pipeline\out\flows\HOOPS_Embedding_Training\flowtrainer\HOOPS_AI_train_SIGNAL_embeddings\0716\091416
-----------------------------------------------------------------------------------


INFO:Restoring states from the checkpoint path at c:\Users\LuisSalazar.LY-LS-LEGION\Documents\repos\HOOPS-AI-tutorials\packages\trained_ml_models\ts3d_2M_hoops_embeddings_SIGNAL-preview.ckpt
INFO:LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
INFO:Restored all states from the checkpoint at c:\Users\LuisSalazar.LY-LS-LEGION\Documents\repos\HOOPS-AI-tutorials\packages\trained_ml_models\ts3d_2M_hoops_embeddings_SIGNAL-preview.ckpt

INFO: Starting epoch 8 with LR: [0.0003]

INFO:Metric eval_loss improved. New best score: 0.239

WARNING: No best checkpoints available yet. Check that the model logs 'eval_loss' during validation.
INFO: Starting epoch 9 with LR: [0.0003]

INFO:Metric eval_loss improved by 0.016 >= min_delta = 0.0. New best score: 0.223
INFO:`Trainer.fit` stopped: `max_epochs=10` reached.

INFO: New best -> best.ckpt (eval_loss=0.23907954)
INFO: New best -> best.ckpt (eval_loss=0.22293355)

[6]:

flow_trainer.test(trained_model_path)

INFO:Restoring states from the checkpoint path at c:\Users\LuisSalazar.LY-LS-LEGION\Documents\repos\HOOPS-AI-tutorials\embeddings_pipeline\out\flows\HOOPS_Embedding_Training\flowtrainer\HOOPS_AI_train_SIGNAL_embeddings\0716\091416\best.ckpt
INFO:LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
INFO:Loaded model weights from the checkpoint at c:\Users\LuisSalazar.LY-LS-LEGION\Documents\repos\HOOPS-AI-tutorials\embeddings_pipeline\out\flows\HOOPS_Embedding_Training\flowtrainer\HOOPS_AI_train_SIGNAL_embeddings\0716\091416\best.ckpt


-----------------------------------------------------------------------------------
HOOPS Embedding Model - TESTING STEP
-----------------------------------------------------------------------------------

The trained model: c:\Users\LuisSalazar.LY-LS-LEGION\Documents\repos\HOOPS-AI-tutorials\embeddings_pipeline\out\flows\HOOPS_Embedding_Training\flowtrainer\HOOPS_AI_train_SIGNAL_embeddings\0716\091416\best.ckpt

Test set contains 12175 training samples
-----------------------------------------------------------------------------------


────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
            Test metric                         DataLoader 0
────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
 test/duplicate_detail_l1_threshold        0.0010000099427998066
test/duplicate_masked_detail_l1_mean                0.0
    test/duplicate_pairs_ignored            0.03090072236955166
          test/loss_curve                   0.07687966525554657
          test/loss_fused                   0.05569160729646683
         test/loss_surface                  0.07976454496383667
          test/loss_total                   0.21233581006526947
────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────

[6]:

{'test/duplicate_pairs_ignored': 0.03090072236955166,
 'test/duplicate_detail_l1_threshold': 0.0010000099427998066,
 'test/duplicate_masked_detail_l1_mean': 0.0,
 'test/loss_total': 0.21233581006526947,
 'test/loss_fused': 0.05569160729646683,
 'test/loss_surface': 0.07976454496383667,
 'test/loss_curve': 0.07687966525554657}

Next steps

Your checkpoint is ready to use anywhere the pre-trained model is used:

from hoops_ai.ml.embeddings import HOOPSEmbeddings

HOOPSEmbeddings.register_model(
    model_name="My Custom Model",
    checkpoint_path=str(trained_model_path),
)

  • Build a similarity index — see `demo_HOOPS_Embeddings_indexing.ipynb <./demo_HOOPS_Embeddings_indexing.ipynb>`__ to embed your full dataset and write a FAISS index with your custom model.

  • Run CAD similarity search — see `5b_cad_search_using_HOOPS_embeddings.ipynb <../notebooks/5b_cad_search_using_HOOPS_embeddings.ipynb>`__ to query the index and retrieve similar parts.

  • Predict missing PLM metadata — see `demo_HOOPS_Embeddings_+_Context_Layer.ipynb <./demo_HOOPS_Embeddings_+_Context_Layer.ipynb>`__ to combine the index with the Context Layer for metadata inference.