Model#
- class tn4ml.models.Model[source][source]#
Bases:
TensorNetworktn4ml.models.Modelclass models training model of classquimb.tensor.tensor_core.TensorNetwork.- loss#
Loss function. See
tn4ml.metricsfor examples.- Type:
Callable, or None
- strategy#
Strategy for computing gradients.
- Type:
tn4ml.strategy.Strategy
- optimizer#
Type of optimizer matching names of optimizers from optax.
- Type:
str
- learning_rate#
Learning rate for optimizer.
- Type:
float
- train_type#
Type of training: 0 = ‘unsupervised’ or 1 =’supervised’, 2 = ‘target TN’.
- Type:
int
- gradient_transforms#
Sequence of gradient transformations.
- Type:
sequence
- device#
Device for computation, e.g. (‘cpu’, 0) or (‘gpu’, 0). [0] = device name, [1] = device index.
- Type:
(str, int)
- opt_state#
State of optimizer.
- Type:
Any
- __init__()[source][source]#
Initialize
tn4ml.models.Model.- Return type:
None
- save(model_name, dir_name='~', tn=False)[source][source]#
Save
tn4ml.models.Modelto pickle file.- Parameters:
model_name (str) – Name of Model.
dir_name (str) – Directory for saving Model.
tn (bool) – If True, model object is TensorNetwork.
- configure(**kwargs)[source][source]#
Configure model for training with specific parameters.
- Parameters:
kwargs (dict)
keys (Configuration parameters. Supported)
strategy (-)
('global' (Training strategy)
'sweeps'
'local'
'dmrg'
'dmrg-like')
optimizer (-) – Optimization algorithm to use.
loss (-) – Loss function for training.
train_type (-) – Type of training (from
tn4ml.util.TrainingType)learning_rate (-) – Learning rate for optimizer.
gradient_transforms (-) – Sequence of gradient transformations for optax
device (-) – Device for computation, e.g. (‘cpu’, 0) or (‘gpu’, 0). [0] = device name, [1] = device index.
- Return type:
None
Examples
>>> model.configure(strategy='global', optimizer=optax.adam, learning_rate=0.01, loss=tn4ml.metrics.LogQuadNorm, train_type=TrainingType.UNSUPERVISED)
- predict(sample, embedding=None, return_tn=False, normalize=False)[source][source]#
Predict the output of the model.
- Parameters:
sample (
numpy.ndarray) – Input data.embedding (
tn4ml.embeddings.Embedding) – Data embedding function.return_tn (bool) – If True, returns tensor network, otherwise returns data. Useful when you want to vmap over predict function.
normalize (bool)
- Returns:
Output of the model.
- Return type:
quimb.tensor.tensor_core.TensorNetwork
- forward(data, embedding=None, batch_size=64, normalize=False, dtype=<class 'jax.numpy.float64'>, seed=42, alternate_flip=False)[source][source]#
Forward pass of the model.
- Parameters:
data (
jax.numpy.ndarray) – Input data.y_true (
jax.numpy.ndarray) – Target class vector.embedding (
tn4ml.embeddings.Embedding) – Data embedding function.batch_size (int) – Batch size for data processing.
normalize (bool) – If True, the model output is normalized.
dtype (Any) – Data type of input data.
seed (int) – Random seed for data shuffling.
alternate_flip (bool)
- Returns:
Output of the model.
- Return type:
jax.numpy.ndarray
- accuracy(data, y_true=None, embedding=None, batch_size=64, shuffle=False, normalize=False, accuracy_fn=None, dtype=<class 'jax.numpy.float64'>, seed=42, alternate_flip=False)[source][source]#
Calculate accuracy for supervised learning.
- Parameters:
model (
tn4ml.models.Model) – Tensor Network model.data (
numpy.ndarray) – Input data.y_true (
numpy.ndarray) – Target class vector.embedding (
tn4ml.embeddings.Embedding) – Data embedding function.batch_size (int) – Batch size for data processing.
normalize (bool) – If True, the model output is normalized in predict function.
accuracy_fn (Callable) – Function applied to raw model outputs before class labels are extracted. If it returns class scores, the predicted class is selected with argmax; if it returns class labels, those labels are compared directly.
dtype (Any) – Data type of input data.
seed (int) – Random seed for data shuffling.
shuffle (bool)
alternate_flip (bool)
- Return type:
float
- update_tensors(params)[source][source]#
Update tensors of the model with new parameters.
- Parameters:
params (sequence of
jax.numpy.ndarray) – New parameters of the model.sitetags (sequence of str, or default None) – Names of tensors for differentiation (for Sweeping strategy).
- Return type:
None
- compute_entropy(data, embedding)[source][source]#
Compute entropy of the model.
- Parameters:
data (
jax.numpy.ndarray) – Input data.embedding (
tn4ml.embeddings.Embedding) – Data embedding function.
- Returns:
Entropy of the model.
- Return type:
float
- compute_entropy_batch(data, embedding)[source][source]#
Compute entropy of the model for a batch of data.
- Parameters:
data (
jax.numpy.ndarray) – Input data.embedding (
tn4ml.embeddings.Embedding) – Data embedding function.
- Returns:
Entropy of the model.
- Return type:
float
- create_train_step(params, loss_func)[source][source]#
Create functions for training steps and gradients.
Creates function for calculating value and gradients of loss, and function for one step in training procedure. Initializes the optimizer and creates optimizer state.
- Parameters:
params (sequence of
jax.numpy.ndarray) – Parameters of the model.loss_func (function) – Loss function.
grads_func (function) – Function for calculating gradients of loss.
- Returns:
train_step (function) – Function to perform one training step.
opt_state (tuple) – State of optimizer at the initialization.
- train(inputs=None, val_inputs=None, targets=None, val_targets=None, tn_target=None, batch_size=None, epochs=1, embedding=None, normalize=False, canonize=(False, None), time_limit=None, earlystop=None, gradient_clip_threshold=None, val_batch_size=None, eval_metric=None, display_val_acc=False, accuracy_fn=None, dtype=<class 'jax.numpy.float64'>, shuffle=False, seed=42, alternate_flip=False)[source][source]#
Perform the training procedure of
tn4ml.models.Model.- Parameters:
inputs (sequence of
numpy.ndarray) – Data used for training procedure.val_inputs (sequence of
numpy.ndarray) – Data used for validation.targets (sequence of
numpy.ndarray) – Targets for training procedure (if training is supervised).val_targets (sequence of
numpy.ndarray) – Targets for validation (if training is supervised).tn_target (
quimb.tensor.tensor_core.TensorNetworkor any specialized TN class from quimb.tensor module) – Target tensor network for training.batch_size (int, or default None) – Number of samples per gradient update.
epochs (int) – Number of epochs for training.
embedding (
tn4ml.embeddings.Embedding) – Data embedding function.normalize (bool) – If True, the model is normalized after each iteration.
canonize (tuple([bool, int])) – tuple indicating is model canonized after each iteration. Example: (True, 0) - model is canonized in canonization center = 0.
time_limit (int) – Time limit on model’s training in seconds.
earlystop (
tn4ml.util.EarlyStopping) – Early stopping training when monitored metric stopped improving.gradient_clip_threshold (float) – Threshold for gradient clipping.
val_batch_size (int) – Number of samples per validation batch.
display_val_acc (bool) – If True, displays validation accuracy.
accuracy_fn (Callable) – Function applied to raw model outputs before validation accuracy labels are extracted. Passed to
accuracy().alternate_flip (bool) – If True, flips every other batch along axis=1.
eval_metric (Callable | None)
dtype (Any)
shuffle (bool | None)
seed (int | None)
- Returns:
history – Records training loss and metric values.
- Return type:
dict
- evaluate(inputs=None, targets=None, tn_target=None, batch_size=None, embedding=None, evaluate_type=TrainingType.UNSUPERVISED, return_list=False, metric=None, dtype=<class 'jax.numpy.float64'>, shuffle=False, seed=42, alternate_flip=False)[source][source]#
Evaluate the model on the data.
- Parameters:
inputs (sequence of
numpy.ndarray) – Data used for evaluation.targets (sequence of
numpy.ndarray) – Targets for evaluation (if evaluation is supervised).tn_target (
quimb.tensor.tensor_core.TensorNetworkor any specialized TN class from quimb) – Target tensor network for evaluation.batch_size (int, or default None) – Number of samples per evaluation.
embedding (
tn4ml.embeddings.Embedding) – Data embedding function.evaluate_type (int) – Type of evaluation: 0 = ‘unsupervised’ or 1 =’unsupervised’.
return_list (bool) – If True, returns list of loss values for each batch.
metric (function) – Metric function for evaluation.
dtype (Any) – Data type of input data.
shuffle (bool) – If True, data is shuffled.
seed (int) – Random seed for data shuffling.
alternate_flip (bool | None)
- Returns:
Loss value.
- Return type:
float
- tn4ml.models.load_model(model_name, dir_name=None)[source][source]#
Load the Model from pickle file.
- Parameters:
model_name (str) – Name of the model.
dir_name (str) – Directory where model is stored.
- Return type:
tn4ml.models.Modelor subclass
- class tn4ml.models.smpo.SpacedMatrixProductOperator[source][source]#
Bases:
TensorNetwork1DOperator,TensorNetwork1DFlat,ModelA MatrixProductOperator with a decimated number of output indices.
See
quimb.tensor.tensor_1d.MatrixProductOperatorfor explanation of other attributes and methods.- __init__(arrays, output_inds=None, shape='lrud', site_tag_id='I{}', tags=None, upper_ind_id='k{}', lower_ind_id='b{}', bond_name='bond{}', **tn_opts)[source][source]#
Create a MatrixProductOperator with a decimated number of output indices.
- Return type:
None
- arrays#
The arrays defining the operator.
- Type:
tuple of array_like
- output_inds#
Indexes of tensors which have output indices. From 0 to n. If spacing is not even.
- Type:
array of int
- shape#
The shape of the tensors, e.g. ‘lurd’ or ‘lrud’.
- Type:
str
- site_tag_id#
The format string for the site tags.
- Type:
str
- tags#
Global tags to add to all tensors.
- Type:
str or sequence of str
- upper_ind_id#
The format string for the upper virtual indices.
- Type:
str
- lower_ind_id#
The format string for the lower virtual indices.
- Type:
str
- bond_name#
The format string for the bond names.
- Type:
str
- tn_opts#
Supplied to
quimb.tensor.tensor_core.TensorNetwork.- Type:
optional
- normalize(insert=None, output_inds=None)[source][source]#
Normalize tensors of
tn4ml.models.smpo.SpacedMatrixProductOperator.- Parameters:
insert (int) – Index of tensor divided by norm. Default = None. When None the norm division is distributed across all tensors.
- Return type:
None
- norm(**contract_opts)[source][source]#
Calculate norm of
tn4ml.models.smpo.SpacedMatrixProductOperator.- Parameters:
contract_opts (Optional) – Arguments passed to
contract().- Returns:
- Return type:
float
- property spacing: int#
Spacing paramater, or space between output indices in number of sites.
- property spacings: list#
Spacings paramater, or space between output indices in number of sites.
- property lower_inds#
Get lower indices for output sites.
- apply_mps(tn_vec, normalize_on_contract=True, compress=False, **compress_opts)[source][source]#
Version of
quimb.tensor.tensor_1d.MatrixProductOperator._apply_mps()fortn4ml.models.smpo.SpacedMatrixProductOperator.- Parameters:
tn_op (
quimb.tensor.tensor_core.TensorNetwork) – The tensor network representing the operator.tn_vec (
quimb.tensor.tensor_core.TensorNetwork, orquimb.tensor.tensor_1d.MatrixProductState) – The tensor network representing the vector.compress (bool) – Whether to compress the resulting tensor network.
compress_opts (optional) – Options to pass to
tn_vec.compress.
- Return type:
quimb.tensor.tensor_1d.MatrixProductState
- apply_smpo(tn_op_2, trace=True, compress=False, **compress_opts)[source][source]#
Version of
quimb.tensor.tensor_1d.MatrixProductOperator._apply_mpo()fortn4ml.models.smpo.SpacedMatrixProductOperator.- Parameters:
tn_op_1 (
tn4ml.models.smpo.SpacedMatrixProductOperator) – The tensor network representing the operator 1.tn_op_2 (
tn4ml.models.smpo.SpacedMatrixProductOperator) – The tensor network representing the operator 2.compress (bool) – Whether to compress the resulting tensor network.
compress_opts (optional) – Options to pass to
tn_vec.compress.
- Return type:
quimb.tensor.tensor_1d.TensorNetwork1D
- apply(other, normalize_on_contract=False, compress=False, **compress_opts)[source][source]#
Apply this SMPO to another SMPO or MPS.
Version of
quimb.tensor.tensor_1d.MatrixProductOperator.apply()fortn4ml.models.smpo.SpacedMatrixProductOperator. Act with this SMPO on another SMPO or MPS, such that the resulting object has the same tensor network structure/indices as other. For an MPS:For an SMPO:
The resulting TN will have the same structure/indices as other, but probably with larger bonds (depending on compression).
- Parameters:
other (
tn4ml.models.smpo.SpacedMatrixProductOperator, orquimb.tensor.tensor_1d.MatrixProductState.) – The object to act on.compress (bool) – Whether to compress the resulting object.
compress_opts (optional) – Supplied to
TensorNetwork1DFlat.compress().
- Return type:
quimb.tensor.tensor_1d.MatrixProductOperator, orquimb.tensor.tensor_1d.MatrixProductState
- class tn4ml.models.mps.MatrixProductState[source][source]#
Bases:
Model,MatrixProductStateA Trainable MatrixProductState class.
See
quimb.tensor.tensor_1d.MatrixProductStatefor explanation of other attributes and methods.- __init__(arrays, **kwargs)[source][source]#
Initialize the MatrixProductState.
- Parameters:
arrays (list of array_like) – The list of tensors, each of shape
(D, D, d), whereDis the bond dimension anddis the physical dimension.**kwargs (dict) – Additional arguments to be passed to the parent class.
- normalize(insert=None)[source][source]#
Normalize tensors of
tn4ml.models.mps.MatrixProductState.
- class tn4ml.models.mpo.MatrixProductOperator[source][source]#
Bases:
Model,MatrixProductOperatorA Trainable MatrixProductOperator class.
See
quimb.tensor.tensor_1d.MatrixProductOperatorfor explanation of other attributes and methods.- __init__(arrays, **kwargs)[source][source]#
Initialize
tn4ml.models.Model.
- normalize(insert=None)[source][source]#
Normalize tensors of
tn4ml.models.mpo.MatrixProductOperator.- Parameters:
insert (int) – Index of tensor divided by norm. Default = None. When None the norm division is distributed across all tensors.
- class tn4ml.models.tn.TensorNetwork[source][source]#
Bases:
Model,TensorNetwork1DFlatA Trainable TensorNetwork class.
See
quimb.tensor.tensor_core.TensorNetworkfor explanation of other attributes and methods.- __init__(tensors, site_tag_id='I{}', cyclic=False, **kwargs)[source][source]#
Initialize
tn4ml.models.tn.ParametrizedTensorNetwork.- Parameters:
tensors (list or TensorNetwork) – List of tensors of
quimb.tensor.tensor_core.Tensoror :class:quimb.tensor.tensor_core.TensorNetwork.kwargs (dict) – Additional arguments.
site_tag_id (str)
cyclic (bool)
- canonize(where, cur_orthog='calc', info=None, bra=None, inplace=False)[source][source]#
Canonizes the tensor network.
- copy(virtual=False, deep=False)[source][source]#
Copy the model.
- Return type:
Model of the same type.
- Parameters:
virtual (bool)
deep (bool)
- norm(**contract_opts)[source][source]#
Calculate norm of
tn4ml.models.tn.TensorNetwork.- Parameters:
contract_opts (Optional) – Arguments passed to
contract().- Returns:
- Return type:
float
- normalize(insert=None)[source][source]#
Normalize tensors of
tn4ml.models.tn.TensorNetwork.- Parameters:
insert (int) – Index of tensor divided by norm. Default = None. When None the norm division is distributed across all tensors.
- Return type:
None