End-to-end training, evaluation, and zkSNARK proof orchestration for zkSTAR.
Manages the full lifecycle of the KFLSTM model: loading config and data,
training the EKF-LSTM sub-models, running the evaluation loop over the test
set, orchestrating ezkl circuit setup and proof generation for both the
prediction and hypothesis circuits, and logging results.
Key environment variables (all optional with sensible defaults):
D_VAL (int): DVAL interval — number of EKF steps accumulated per proof. Default 120.
HYP_TEST_INT (int): Hypothesis test interval in timesteps. Default 128.
GEN_PREDICTION_PROOFS (int): 0/1 flag to disable/enable prediction proof generation.
GEN_HYPOTHESIS_PROOFS (int): 0/1 flag to disable/enable hypothesis proof generation.
PRELOADED_PROOF_ROOT (str): Path to pre-existing VK/PK/settings artefacts.
If set and the path exists, skips the ezkl setup phase.
ALPHA (float): Significance level for the chi-squared hypothesis test. Default 0.005.
NUM_THREADS (int): Number of PyTorch CPU threads.
TRAIN_RATIO (float): Fraction of data used for training. Default 0.8.
WRITE_WEIGHT_FILE (str): Path to save trained model weights.
VECTORIZED (int): 0 to disable vectorized EKF covariance updates.
DIAG_COV (int): If set, use diagonal-only covariance approximation.
Parameters:
| Name |
Type |
Description |
Default |
config_file_name
|
str
|
Path to the YAML configuration file specifying
dataset paths, sensor/state dictionaries, model hyperparameters, and
result output directories.
|
required
|
write_weight_file
|
str
|
Override path for saving model weights.
If None, falls back to WRITE_WEIGHT_FILE env var or a UUID-based filename.
|
None
|
Source code in src/KFLSTMTrainer_ZKL.py
| def __init__(self,config_file_name,write_weight_file=None):
self.exp_id = str(uuid.uuid4())
if "NUM_THREADS" in os.environ.keys():
torch.set_num_threads(int(os.environ["NUM_THREADS"]))
self.lr = 1e-3
self.la_mse = 1
self.la_nll_r = 1
self.la_nll_h = 1
self.la_CPCT_cov = 1
self.la_K = 1
# self.start_test = 0
# self.end_test = self.test_window
# self.start_forecast = self.test_window
# self.end_forecast = self.test_window+self.forecast_window
self.num_steps = 1
self.look_ahead = 1
self.train_ratio = 0.8
self.outputs = None
self.k1 = 256
if dist_train:
self.comm = MPI.COMM_WORLD
self.rank = self.comm.Get_rank()
self.size = self.comm.Get_size()
else:
self.rank = 0
self.size = 1
self.comm = None
self.format_string = "%Y-%m-%d %H:%M:%S" #%Y-%m-%d %H:%M:%S
self.D = 120
if "D_VAL" in os.environ.keys():
self.D = int(os.environ["D_VAL"])
self.HYPOTHESIS_TEST_INTERVAL = 128
if "HYP_TEST_INT" in os.environ.keys():
self.HYPOTHESIS_TEST_INTERVAL = int(os.environ["HYP_TEST_INT"])
self.timestamped_proof_folder_name=True
if "TIMESTAMPED_PROOF_FOLDER" in os.environ.keys():
self.timestamped_proof_folder_name = int(os.environ["TIMESTAMPED_PROOF_FOLDER"])
self.preloaded_proof_root = None
if "PRELOADED_PROOF_ROOT" in os.environ.keys():
self.preloaded_proof_root = os.environ["PRELOADED_PROOF_ROOT"]
if self.preloaded_proof_root == "none":
self.preloaded_proof_root = None
if self.preloaded_proof_root is None or not os.path.exists(self.preloaded_proof_root):
self.setup_prediction_snark = False
self.setup_hypothesis_snark = False
else:
self.setup_prediction_snark = True
self.setup_hypothesis_snark = True
self.gen_prediction_proofs=True
if "GEN_PREDICTION_PROOFS" in os.environ.keys():
self.gen_prediction_proofs = int(os.environ["GEN_PREDICTION_PROOFS"])
self.gen_hypothesis_proofs=True
if "GEN_HYPOTHESIS_PROOFS" in os.environ.keys():
self.gen_hypothesis_proofs = int(os.environ["GEN_HYPOTHESIS_PROOFS"])
self.initial_time_stamp = 0
self.start_timestamp = 0
self.duration = 0
self.attack_timestamp = 0
self.attack_id = "normal"
self.retain_graph = True
self.train_mode = True
self.processConfig(config_file_name)
self.proof_path = self.result_folder_name
self.shuffle = False
self.vectorized = True
if "VECTORIZED" in os.environ.keys() and not int(os.environ["VECTORIZED"]):
self.vectorized= False
self.diag_cov = 0
self.minmaxscaler = True
self.h_out_unroll = False
if "HUNROLL" in os.environ.keys():
self.h_out_unroll = True
if "DIAG_COV" in os.environ.keys():
self.diag_cov = int(os.environ["DIAG_COV"])
if "SHUFFLE" in os.environ.keys():
self.shuffle = True
if "NUM_THREADS" in os.environ.keys():
torch.set_num_threads(int(os.environ["NUM_THREADS"]))
if "N_STEPS" in os.environ.keys():
self.num_steps = int(os.environ["N_STEPS"])
if "TRAIN_RATIO" in os.environ.keys():
self.train_ratio = float(os.environ["TRAIN_RATIO"])
if write_weight_file is not None:
self.write_weight_file = write_weight_file
elif "WRITE_WEIGHT_FILE" in os.environ.keys():
self.write_weight_file = os.environ["WRITE_WEIGHT_FILE"]
else:
self.write_weight_file = self.exp_id+".pth"
if "ALPHA" in os.environ.keys():
self.alpha = float(os.environ["ALPHA"])
self.original_alpha = float(os.environ["ALPHA"])
self.perturbed_alpha = []
else:
self.alpha=0.005
self.write_weight_file = os.path.join(self.result_root,self.write_weight_file)
print(f"Weights being written to {self.write_weight_file}")
self.initData(self.state_dict, self.measure_dict, self.data_file_name)
self.prepareData()
self.slice_len = len(self.train_data)
self.local_slices = [0]
self.dist_comm = dist_comm
# Prepare empty tensors to hold the summed parameters on each rank
self.grad_sum = None
self.output_ezkl = None
self.chi2_ul = 100
if "CHI2_UL" in os.environ.keys():
self.chi2_ul = float(os.environ["CHI2_UL"])
self.chi2_num_points = self.chi2_ul
if "CHI2_NUM_POINTS" in os.environ.keys():
self.chi2_num_points = int(os.environ["CHI2_NUM_POINTS"])
self.zkp_scale = 12
if "ZKSNARK_SCALE" in os.environ.keys():
self.zkp_scale = int(os.environ["ZKSNARK_SCALE"])
self.prediction_setup_times_files = "prediction_setup_times_files.csv"
self.hypothesis_setup_times_files = "hypothesis_setup_times_files.csv"
self.prediction_prove_times_files = "prediction_prove_times_files.csv"
self.hypothesis_prove_times_files = "hypothesis_prove_times_files.csv"
self.svd_check_files = "svd_check_files.json"
self.kcr_check_files = "kcr_check_files.json"
|
dump_yaml_data
dump_yaml_data(data, yaml_file)
Dumps data to a YAML file.
Parameters:
| Name |
Type |
Description |
Default |
data
|
|
|
required
|
yaml_file
|
|
The path to the YAML file.
|
required
|
Source code in src/KFLSTMTrainer_ZKL.py
| def dump_yaml_data(self,data, yaml_file):
"""Dumps data to a YAML file.
Args:
data: The data to be dumped.
yaml_file: The path to the YAML file.
"""
with open(yaml_file, 'w') as f:
yaml.dump(data, f, default_flow_style=False)
|
read_yaml_data
read_yaml_data(yaml_file)
Reads data from a YAML file.
Parameters:
| Name |
Type |
Description |
Default |
yaml_file
|
|
The path to the YAML file.
|
required
|
Returns:
| Type |
Description |
|
|
A dictionary containing the data from the YAML file.
|
Source code in src/KFLSTMTrainer_ZKL.py
| def read_yaml_data(self,yaml_file):
"""Reads data from a YAML file.
Args:
yaml_file: The path to the YAML file.
Returns:
A dictionary containing the data from the YAML file.
"""
with open(yaml_file, 'r') as f:
data = yaml.safe_load(f)
return data
|