Sweeping method for TN training

Classification of Breast Cancer dataset using Sweeping method

import os
from pathlib import Path

os.environ["KMP_WARNINGS"] = "0"

import jax
import jax.numpy as jnp
import numpy as np
import optax
import pandas as pd
from jax.nn.initializers import *
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.utils.class_weight import compute_class_weight

from tn4ml.embeddings import *
from tn4ml.eval import *
from tn4ml.initializers import *
from tn4ml.metrics import *
from tn4ml.models.model import *
from tn4ml.models.mps import *
from tn4ml.util import *

# Enable 64-bit precision and set matmul precision to highest
jax.config.update("jax_enable_x64", True)
jax.config.update("jax_default_matmul_precision", "highest")

Load dataset

The dataset can be download from kaggle.com/breast-cancer/data.

load_dir = "data"
save_dir = "results"
n_classes = 2

# Load data
breast_cancer_path = Path(load_dir) / "breast-cancer.csv"
if breast_cancer_path.exists():
    data = pd.read_csv(breast_cancer_path)
else:
    breast_cancer = load_breast_cancer(as_frame=True)
    data = breast_cancer.frame.copy()
    data.insert(0, "id", np.arange(len(data)))
    data["diagnosis"] = data["target"].map({1: "B", 0: "M"})
    data = data.drop(columns=["target"])

data["diagnosis"] = data["diagnosis"].map({"M": 1, "B": 0})

X = data.drop(["id", "diagnosis"], axis=1)
y = data["diagnosis"].to_numpy()

feature_min = X.min()
feature_max = X.max()

X_normalized = (X - feature_min) / (feature_max - feature_min)
X_numpy = X_normalized.to_numpy()

if os.environ.get("CI"):
    X_numpy = X_numpy[:, :2]

X_train, X_test, y_train, y_test = train_test_split(
    X_numpy, y, test_size=0.2, random_state=42
)
X_train, X_valid, y_train, y_valid = train_test_split(
    X_train, y_train, test_size=0.2, random_state=42
)

if os.environ.get("CI"):
    X_train, y_train = X_train[:4], y_train[:4]
    X_valid, y_valid = X_valid[:4], y_valid[:4]
    X_test, y_test = X_test[:4], y_test[:4]

print("Train data shape: ", X_train.shape)
print("Validation data shape: ", X_valid.shape)
print("Test data shape: ", X_test.shape)

classes = np.unique(y_train)
print("Classes: ", classes)
class_weights = compute_class_weight(
    class_weight="balanced", classes=classes, y=y_train
)
print("Class weights: ", class_weights)

y_train = integer_to_one_hot(y_train, n_classes)
y_valid = integer_to_one_hot(y_valid, n_classes)
y_test = integer_to_one_hot(y_test, n_classes)

Train data shape:  (364, 30)
Validation data shape:  (91, 30)
Test data shape:  (114, 30)
Classes:  [0 1]
Class weights:  [0.78787879 1.36842105]

Training setup

Define model parameters

key = jax.random.key(42)
L = X_train.shape[1]
print("Number of tensors: ", L)
class_index = int(L // 2)
shape_method = "noteven"  # default method
compress = False  # connected with shape method
embedding = PolynomialEmbedding(
    degree=2, n=1, include_bias=True
)  # polynomial embedding of degree 2
phys_dim = 3
initializer = randn(0.9)
bond_dim = 1 if os.environ.get("CI") else 10

Number of tensors:  30

Initialize the MPS model

model = MPS_initialize(
    L=L,
    initializer=initializer,
    key=key,
    shape_method=shape_method,
    compress=compress,
    cyclic=False,
    phys_dim=phys_dim,
    bond_dim=bond_dim,
    class_index=class_index,
    canonical_center=class_index,
    class_dim=n_classes,
    add_identity=True,
    boundary="obc",
)

Define training parameters

def weighted_crossentropy_loss(*args, **kwargs):
    """Compute weighted cross-entropy loss."""
    return CrossEntropyWeighted(class_weights=class_weights)(*args, **kwargs).mean()


def crossentropy_loss(*args, **kwargs):
    """Compute softmax cross-entropy loss."""
    return OptaxWrapper(optax.softmax_cross_entropy)(*args, **kwargs).mean()

learning_rate = 1e-4
optimizer = optax.adam
strategy = "global" if os.environ.get("CI") else "sweeps"
loss = weighted_crossentropy_loss
train_type = TrainingType.SUPERVISED
earlystop = EarlyStopping(min_delta=0, patience=5, monitor="loss", mode="min")
epochs = 1 if os.environ.get("CI") else 20
batch_size = 2 if os.environ.get("CI") else 128
val_batch_size = 2 if os.environ.get("CI") else 30
train_dtype = jnp.float32 if os.environ.get("CI") else jnp.float64

model.configure(
    optimizer=optimizer,
    strategy=strategy,
    loss=loss,
    train_type=train_type,
    learning_rate=learning_rate,
)

history = model.train(
    X_train,
    targets=y_train,
    val_inputs=X_valid,
    val_targets=y_valid,
    epochs=epochs,
    batch_size=batch_size,
    embedding=embedding,
    normalize=True,
    dtype=train_dtype,
    earlystop=earlystop,
    canonize=(True, class_index),
    display_val_acc=not bool(os.environ.get("CI")),
    eval_metric=crossentropy_loss,
    val_batch_size=val_batch_size,
)

epoch: 100%|██████████ 20/20 , loss=0.8486, val_loss=0.3949, val_acc=0.8889

Plot loss

plot_loss(history, validation=True, figsize=(8, 6))