Training an Imported ONNX Model¶

In this example, we load an existing ONNX network and a VNNLIB specification, then train the network towards the desired property while preserving its original outputs on sampled inputs. Supply your own network.onnx and property.vnnlib files in the notebook's working directory. The example assumes vector outputs suitable for a mean-squared-error task loss.

First, we import the required functions.

In [ ]:

import copy
import torch
from torch.utils.data import DataLoader, TensorDataset

from CTRAIN.data_loaders import load_vnnlib_problem
from CTRAIN.model_wrappers import PropertyModelWrapper

import copy import torch from torch.utils.data import DataLoader, TensorDataset from CTRAIN.data_loaders import load_vnnlib_problem from CTRAIN.model_wrappers import PropertyModelWrapper

Now, we load the network and specification. By default, VNNLIB output assertions describe unsafe states; CTRAIN complements them to obtain the desired property. Input coordinates are used as declared, without inferred normalisation.

In [ ]:

model, property_spec, provenance = load_vnnlib_problem(
    "network.onnx", "property.vnnlib",
)
input_shape = tuple(property_spec.domains[0].lower.shape)

model, property_spec, provenance = load_vnnlib_problem( "network.onnx", "property.vnnlib", ) input_shape = tuple(property_spec.domains[0].lower.shape)

We sample task inputs from the first declared box and use the original model outputs as preservation targets. For an application, replace these samples with representative task data.

In [ ]:

reference_model = copy.deepcopy(model).eval()
box = property_spec.domains[0]
parameter = next(model.parameters())
inputs = (box.lower + torch.rand(256, *input_shape) * (box.upper - box.lower)).to(parameter)
with torch.no_grad():
    targets = reference_model(inputs)
train_loader = DataLoader(TensorDataset(inputs, targets), batch_size=32, shuffle=True)

reference_model = copy.deepcopy(model).eval() box = property_spec.domains[0] parameter = next(model.parameters()) inputs = (box.lower + torch.rand(256, *input_shape) * (box.upper - box.lower)).to(parameter) with torch.no_grad(): targets = reference_model(inputs) train_loader = DataLoader(TensorDataset(inputs, targets), batch_size=32, shuffle=True)

Now, we wrap the imported model and compute its initial certification result. All declared boxes participate in property training and certification, even though the task samples above use only the first box.

In [ ]:

wrapped_model = PropertyModelWrapper(
    model=model,
    input_shape=input_shape,
    properties=[property_spec],
    task_loss=torch.nn.MSELoss(),
    num_epochs=20,
    training_method="CROWN-IBP",
    property_weight=10.0,
    device="cpu",
)
before = wrapped_model.certify_properties(method="CROWN")

wrapped_model = PropertyModelWrapper( model=model, input_shape=input_shape, properties=[property_spec], task_loss=torch.nn.MSELoss(), num_epochs=20, training_method="CROWN-IBP", property_weight=10.0, device="cpu", ) before = wrapped_model.certify_properties(method="CROWN")

Finally, we train and compute fresh bounds for the resulting model.

In [ ]:

wrapped_model.train_model(train_loader)
after = wrapped_model.certify_properties(method="CROWN")
print("Certified before training:", before["certified"])
print("Certified after training:", after["certified"])

wrapped_model.train_model(train_loader) after = wrapped_model.certify_properties(method="CROWN") print("Certified before training:", before["certified"]) print("Certified after training:", after["certified"])

The task loss encourages preservation on the sampled inputs; it does not certify unchanged behaviour elsewhere. Training may leave the property unresolved. The property guide describes supported import constructs, safe-output assertions and complete verification.