Examples#
Scripts in the examples/ folder of the repository. Run python <script> --help for all options.
Basic#
xor.py#
Train discrete or guided TM on noisy XOR (parity of –nbits relevant bits, plus –noisybits random bits).
"""Train discrete or guided TM on noisy XOR (parity of --nbits relevant bits, plus --noisybits random bits).
Training labels are flipped with probability --noise. Test set is clean, so 100% test accuracy means every xor pattern was learned.
Usage:
python xor.py {discrete,guided} [options]
Examples:
python xor.py discrete # 2-bit XOR, 10% label noise
python xor.py discrete --noisybits 10 # classic noisy XOR: 2 relevant + 10 random bits
python xor.py discrete --nbits 3 --n_clauses 8 --s 8 --T 4 --save xor3.ytm
python xor.py guided --noise 0.0
Options:
- python xor.py <discrete/guided> --help shows all options.
- Common: --nbits --noisybits --noise --N --repeat_test --epochs --n_clauses --s --n_states --weighted 0|1 --coalesced 0|1
--boost_tp_inc 0|1 --boost_tp_dec 0|1 --seed --device cpu:N|cuda:N --save PATH --print
- discrete: --T
- guided: --lr --lambda
"""
import argparse
import pickle
import numpy as np
from ytm.guided import BinaryTM as GuidedTM
from ytm.tm import BinaryTM as DiscreteTM
def make_xor(N: int, nbits: int, noisybits: int, noise: float, rng: np.random.Generator):
X = rng.integers(0, 2, size=(N, nbits + noisybits), dtype=np.uint8)
Y = X[:, :nbits].sum(axis=1) % 2
flip = rng.random(N) < noise
Y[flip] = 1 - Y[flip]
return X, Y.astype(np.uint8)
def make_test(nbits: int, noisybits: int, repeat: int, rng: np.random.Generator):
# All 2**nbits relevant patterns, each repeated with random noisy bits, clean labels.
patterns = ((np.arange(2**nbits)[:, None] >> np.arange(nbits)) & 1).astype(np.uint8)
X_rel = np.repeat(patterns, repeat, axis=0)
X = np.concatenate([X_rel, rng.integers(0, 2, size=(len(X_rel), noisybits), dtype=np.uint8)], axis=1)
return X, (X_rel.sum(axis=1) % 2).astype(np.uint8)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="TM noisy XOR")
# Subcommand to select discrete or guided
models = parser.add_subparsers(dest="model", required=True)
# Common args
common = argparse.ArgumentParser(add_help=False)
common.add_argument("--N", type=int, default=1000, help="training set size")
common.add_argument("--repeat_test", type=int, default=1, help="test set = all 2^nbits patterns, each repeated this many times")
common.add_argument("--nbits", type=int, default=2, help="relevant bits, label is their parity")
common.add_argument("--noisybits", type=int, default=0, help="extra random bits")
common.add_argument("--noise", type=float, default=0.1, help="add noise to training labels")
common.add_argument("--epochs", type=int, default=1)
common.add_argument("--n_clauses", type=int, default=4)
common.add_argument("--s", type=float, default=4.0)
common.add_argument("--n_states", type=int, default=10)
common.add_argument("--weighted", type=lambda v: bool(int(v)), default=True, metavar="0|1")
common.add_argument("--coalesced", type=lambda v: bool(int(v)), default=False, metavar="0|1")
common.add_argument("--boost_tp_inc", type=lambda v: bool(int(v)), default=False, metavar="0|1")
common.add_argument("--boost_tp_dec", type=lambda v: bool(int(v)), default=False, metavar="0|1")
common.add_argument("--seed", type=int, default=10)
common.add_argument("--device", type=str, default="cpu:1", help="cpu:N_THREADS or cuda:GPU_ID")
common.add_argument("--save", type=str, default=None, help="pickle the trained model to this path, e.g. xor.ytm")
common.add_argument("--print", dest="print_clauses", action="store_true", help="print the learned clauses after training")
# Discrete args
discrete = models.add_parser("discrete", parents=[common])
discrete.add_argument("--T", type=float, default=2)
# Guided args
guided = models.add_parser("guided", parents=[common])
guided.add_argument("--lr", type=float, default=1.0)
guided.add_argument("--lambda", dest="lambda_", type=float, default=1.0)
# Parse args
args = parser.parse_args()
# Args to dict
params = vars(args)
model_name, epochs, save = params.pop("model"), params.pop("epochs"), params.pop("save")
nbits, noisybits, noise, N = params.pop("nbits"), params.pop("noisybits"), params.pop("noise"), params.pop("N")
repeat_test, show_clauses = params.pop("repeat_test"), params.pop("print_clauses")
# Data: noisy training labels, clean test labels
rng = np.random.default_rng(params["seed"] + 50)
X_train, Y_train = make_xor(N, nbits, noisybits, noise, rng)
X_test, Y_test = make_test(nbits, noisybits, repeat_test, rng)
# Model Initialization
TM = DiscreteTM if model_name == "discrete" else GuidedTM
tm = TM(**params, dim=(nbits + noisybits, 1, 1))
# Training
for epoch in range(epochs):
tm.fit(X_train, Y_train)
train_acc = (tm.predict(X_train)[0].reshape(-1) == Y_train).mean()
test_acc = (tm.predict(X_test)[0].reshape(-1) == Y_test).mean()
print(f"{model_name} epoch {epoch + 1}/{epochs} train acc {train_acc * 100:.2f}% test acc {test_acc * 100:.2f}%")
if show_clauses:
# Relevant bits are x0, x1, ..., the random ones n0, n1, ...
tm.get_clauses().print([f"x{i}" for i in range(nbits)] + [f"n{i}" for i in range(noisybits)], sort="w0,len")
if save is not None:
with open(save, "wb") as f:
pickle.dump(tm, f)
print(f"saved to {save}")
# Loading it back (always lands on cpu:1, use .to() to move it):
# with open(save, "rb") as f:
# tm = pickle.load(f)
# tm.to("cuda:0")
# pred, _ = tm.predict(X_test)
iris.py#
Train discrete or guided TM on Iris (4 numeric features, 3 classes).
"""Train discrete or guided TM on Iris (4 numeric features, 3 classes).
Usage:
python iris.py {discrete,guided} [options]
Examples:
python iris.py guided # 2 clauses per class
python iris.py discrete --T 10
python iris.py guided --n_clauses 4 --coalesced 1 # 4 clauses shared by all classes
Options:
- python iris.py <discrete/guided> --help shows all options.
- Common: --epochs --bins --strategy uniform|quantile --n_clauses --coalesced 0|1 --s --max_includes
--boost_tp_inc 0|1 --boost_tp_dec 0|1 --seed --device cpu:N|cuda:N --print
- discrete: --T
- guided: --lr --lambda
"""
import argparse
import numpy as np
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import KBinsDiscretizer
from ytm.guided import MultiClassTM as GuidedTM
from ytm.tm import MultiClassTM as DiscreteTM
from ytm.utils import Timer, print_table
def load(bins: int, strategy: str, seed: int):
X, Y = load_iris(return_X_y=True)
X_train, X_test, Y_train, Y_test = train_test_split(X, Y.astype(np.uint8), test_size=0.3, stratify=Y, random_state=seed)
d = KBinsDiscretizer(n_bins=bins, encode="ordinal", strategy=strategy)
X_train, X_test = d.fit_transform(X_train).astype(np.int32), d.transform(X_test).astype(np.int32)
return X_train, Y_train, X_test, Y_test, d.n_bins_ - 1
def train_model(tm, xtrain, ytrain, xtest, ytest, epochs: int, model_name: str):
for epoch in range(epochs):
with (fit_timer := Timer()):
loss = tm.fit(xtrain, ytrain)
with (test_timer := Timer()):
test_pred, _ = tm.predict(xtest)
with (train_timer := Timer()):
train_pred, _ = tm.predict(xtrain)
train_log = {
"Acc": f"{(train_pred == ytrain).mean() * 100:.2f}%",
"Eval Time": f"{train_timer.elapsed:.2f}s",
"Fit Time": f"{fit_timer.elapsed:.2f}s",
}
test_log = {"Acc": f"{(test_pred == ytest).mean() * 100:.2f}%", "Eval Time": f"{test_timer.elapsed:.2f}s"}
if loss is not None: # Only loss-guided TM returns loss
train_log["Loss"] = f"{loss:.4f}"
print_table(
f"{model_name} epoch {epoch + 1}/{epochs}",
{
"Train": train_log,
"Test": test_log,
},
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="TM Iris dataset")
# Subcommand to select discrete or guided
models = parser.add_subparsers(dest="model", required=True)
# Common args
common = argparse.ArgumentParser(add_help=False)
common.add_argument("--epochs", type=int, default=200)
common.add_argument("--bins", type=int, default=8, help="discretization levels per feature")
common.add_argument("--strategy", choices=["uniform", "quantile"], default="quantile")
common.add_argument("--n_clauses", type=int, default=2, help="per class, or total with --coalesced 1")
common.add_argument("--coalesced", type=lambda v: bool(int(v)), default=False, metavar="0|1")
common.add_argument("--s", type=float, default=1.5)
common.add_argument("--max_includes", type=int, default=4, help="max literals per clause")
common.add_argument("--boost_tp_inc", type=lambda v: bool(int(v)), default=False, metavar="0|1")
common.add_argument("--boost_tp_dec", type=lambda v: bool(int(v)), default=False, metavar="0|1")
common.add_argument("--seed", type=int, default=10)
common.add_argument("--device", type=str, default="cpu:1", help="cpu:N_THREADS or cuda:GPU_ID")
common.add_argument("--print", dest="print_clauses", action="store_true", help="print the learned clauses after training")
# Discrete args
discrete = models.add_parser("discrete", parents=[common])
discrete.add_argument("--T", type=float, default=10)
# Guided args
guided = models.add_parser("guided", parents=[common])
guided.add_argument("--lr", type=float, default=0.03)
guided.add_argument("--lambda", dest="lambda_", type=float, default=1.0)
# Parse args
args = parser.parse_args()
# Args to dict
params = vars(args)
model_name, epochs = params.pop("model"), params.pop("epochs")
bins, strategy, show_clauses = params.pop("bins"), params.pop("strategy"), params.pop("print_clauses")
X_train, Y_train, X_test, Y_test, feat_maxs = load(bins, strategy, params["seed"])
# Model Initialization
TM = DiscreteTM if model_name == "discrete" else GuidedTM
tm = TM(**params, dim=(X_train.shape[1], 1, 1), n_classes=3, feat_maxs=feat_maxs)
# Training
train_model(tm, X_train, Y_train, X_test, Y_test, epochs=epochs, model_name=model_name)
if show_clauses:
# Bounds are in bin units, 0 .. bins - 1
tm.get_clauses().print(["sepal_len", "sepal_wid", "petal_len", "petal_wid"], sort="w0,len")
Image#
image/mnist.py#
Train discrete or guided TM on MNIST dataset.
"""Train discrete or guided TM on MNIST dataset.
Usage:
python mnist.py {discrete,guided} [options]
Examples:
python mnist.py discrete # 10x10 conv, cpu:1
python mnist.py discrete --patch none --T 5000 # flat, no conv
python mnist.py guided --lr 0.5 --device cuda:0 --epochs 20
Options:
- python mnist.py <discrete/guided> --help shows all options.
- Common: --epochs --n_clauses --s --patch H,W|none --seed --device cpu:N|cuda:N
- discrete: --T
- guided: --lr --lambda
"""
import argparse
import numpy as np
from datasets import load_dataset
from ytm.tm import MultiClassTM as DiscreteTM
from ytm.guided import MultiClassTM as GuidedTM
from ytm.utils import Timer, print_table
def load_mnist():
ds = load_dataset("ylecun/mnist")
X_train, X_test = ((np.array(ds[split]["image"]) > 75).astype(np.uint8) for split in ("train", "test"))
Y_train, Y_test = (np.array(ds[split]["label"], dtype=np.uint8) for split in ("train", "test"))
return X_train, Y_train, X_test, Y_test
def train_model(tm, xtrain, ytrain, xtest, ytest, epochs: int, model_name: str):
for epoch in range(epochs):
with (fit_timer := Timer()):
loss = tm.fit(xtrain, ytrain)
with (test_timer := Timer()):
test_pred, _ = tm.predict(xtest)
with (train_timer := Timer()):
train_pred, _ = tm.predict(xtrain)
train_log = {
"Acc": f"{(train_pred == ytrain).mean() * 100:.4f}%",
"Eval Time": f"{train_timer.elapsed:.2f}s",
"Fit Time": f"{fit_timer.elapsed:.2f}s",
}
test_log = {"Acc": f"{(test_pred == ytest).mean() * 100:.4f}%", "Eval Time": f"{test_timer.elapsed:.2f}s"}
if loss is not None: # Only loss-guided TM returns loss
train_log["Loss"] = f"{loss:.4f}"
print_table(
f"{model_name} epoch {epoch + 1}/{epochs}",
{
"Train": train_log,
"Test": test_log,
},
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="TM MNIST dataset")
# Subcommand to select discrete or guided
models = parser.add_subparsers(dest="model", required=True)
# Commmon args
common = argparse.ArgumentParser(add_help=False)
common.add_argument("--epochs", type=int, default=10)
common.add_argument("--n_clauses", type=int, default=500)
common.add_argument("--s", type=float, default=10.0)
common.add_argument(
"--patch",
dest="patch_dim",
type=lambda v: None if v.lower() == "none" else tuple(map(int, v.split(","))),
default="10,10",
help="H,W or none",
)
common.add_argument("--seed", type=int, default=10)
common.add_argument("--device", type=str, default="cpu:1", help="cpu:N_THREADS or cuda:GPU_ID")
# Discrete args
discrete = models.add_parser("discrete", parents=[common])
discrete.add_argument("--T", type=float, default=1000)
# Guided args
guided = models.add_parser("guided", parents=[common])
guided.add_argument("--lr", type=float, default=1.0)
guided.add_argument("--lambda", dest="lambda_", type=float, default=1.0)
# Parse args
args = parser.parse_args()
# Args to dict
params = vars(args)
model_name, epochs = params.pop("model"), params.pop("epochs")
# Model Initialization
TM = DiscreteTM if model_name == "discrete" else GuidedTM
tm = TM(
**params,
dim=(28, 28, 1),
n_classes=10,
)
# Training
train_model(tm, *load_mnist(), epochs=epochs, model_name=model_name)
image/fmnist.py#
Train discrete or guided TM on Fashion-MNIST dataset.
"""Train discrete or guided TM on Fashion-MNIST dataset.
Usage:
python fmnist.py {discrete,guided} [options]
Examples:
python fmnist.py discrete # 3x3 conv, cpu:1
python fmnist.py discrete --patch none --T 5000 # flat, no conv
python fmnist.py guided --lr 0.5 --device cuda:0 --epochs 20
Options:
- python fmnist.py <discrete/guided> --help shows all options.
- Common: --epochs --clause_drop_p --levels --n_clauses --s --patch H,W|none --seed --device cpu:N|cuda:N
- discrete: --T
- guided: --lr --lambda
"""
import argparse
import numpy as np
from datasets import load_dataset
from ytm.tm import MultiClassTM as DiscreteTM
from ytm.guided import MultiClassTM as GuidedTM
from ytm.utils import Timer, print_table
def load_fmnist(levels: int):
ds = load_dataset("zalando-datasets/fashion_mnist")
# Pixels 0-255 -> levels 0-levels, thermometer encoded by the TM.
X_train, X_test = ((levels * np.array(ds[split]["image"], dtype=np.float32) / 255.0).astype(np.int32) for split in ("train", "test"))
Y_train, Y_test = (np.array(ds[split]["label"], dtype=np.uint8) for split in ("train", "test"))
return X_train, Y_train, X_test, Y_test
def train_model(tm, xtrain, ytrain, xtest, ytest, epochs: int, clause_drop_p: float, model_name: str):
for epoch in range(epochs):
with (fit_timer := Timer()):
loss = tm.fit(xtrain, ytrain, clause_drop_p=clause_drop_p)
with (test_timer := Timer()):
test_pred, _ = tm.predict(xtest)
with (train_timer := Timer()):
train_pred, _ = tm.predict(xtrain)
train_log = {
"Acc": f"{(train_pred == ytrain).mean() * 100:.4f}%",
"Eval Time": f"{train_timer.elapsed:.2f}s",
"Fit Time": f"{fit_timer.elapsed:.2f}s",
}
test_log = {"Acc": f"{(test_pred == ytest).mean() * 100:.4f}%", "Eval Time": f"{test_timer.elapsed:.2f}s"}
if loss is not None: # Only loss-guided TM returns loss
train_log["Loss"] = f"{loss:.4f}"
print_table(
f"{model_name} epoch {epoch + 1}/{epochs}",
{
"Train": train_log,
"Test": test_log,
},
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="TM Fashion-MNIST dataset")
# Subcommand to select discrete or guided
models = parser.add_subparsers(dest="model", required=True)
# Common args
common = argparse.ArgumentParser(add_help=False)
common.add_argument("--epochs", type=int, default=10)
common.add_argument("--clause_drop_p", type=float, default=0.0)
common.add_argument("--levels", type=int, default=8, help="thermometer levels")
common.add_argument("--n_clauses", type=int, default=40000)
common.add_argument("--s", type=float, default=10.0)
common.add_argument(
"--patch",
dest="patch_dim",
type=lambda v: None if v.lower() == "none" else tuple(map(int, v.split(","))),
default="3,3",
help="H,W or none",
)
common.add_argument("--seed", type=int, default=10)
common.add_argument("--device", type=str, default="cpu:1", help="cpu:N_THREADS or cuda:GPU_ID")
# Discrete args
discrete = models.add_parser("discrete", parents=[common])
discrete.add_argument("--T", type=float, default=15000)
# Guided args
guided = models.add_parser("guided", parents=[common])
guided.add_argument("--lr", type=float, default=0.5)
guided.add_argument("--lambda", dest="lambda_", type=float, default=3.0)
# Parse args
args = parser.parse_args()
# Args to dict
params = vars(args)
model_name, epochs, clause_drop_p = params.pop("model"), params.pop("epochs"), params.pop("clause_drop_p")
levels = params.pop("levels")
# Model Initialization
TM = DiscreteTM if model_name == "discrete" else GuidedTM
tm = TM(
**params,
dim=(28, 28, 1),
n_classes=10,
feat_maxs=levels,
)
# Training
train_model(tm, *load_fmnist(levels), epochs=epochs, clause_drop_p=clause_drop_p, model_name=model_name)
image/kmnist.py#
Train discrete or guided TM on Kuzushiji-MNIST dataset.
"""Train discrete or guided TM on Kuzushiji-MNIST dataset.
Usage:
python kmnist.py {discrete,guided} [options]
Examples:
python kmnist.py discrete # 10x10 conv, cpu:1
python kmnist.py discrete --patch none --T 5000 # flat, no conv
python kmnist.py guided --lr 0.5 --device cuda:0 --epochs 20
Options:
- python kmnist.py <discrete/guided> --help shows all options.
- Common: --epochs --clause_drop_p --n_clauses --s --coalesced 0|1 --patch H,W|none --seed --device cpu:N|cuda:N
- discrete: --T
- guided: --lr --lambda
"""
import argparse
import numpy as np
from datasets import load_dataset
from ytm.tm import MultiClassTM as DiscreteTM
from ytm.guided import MultiClassTM as GuidedTM
from ytm.utils import Timer, print_table
def load_kmnist():
ds = load_dataset("tanganke/kmnist")
X_train, X_test = ((np.array(ds[split]["image"]) > 75).astype(np.uint8) for split in ("train", "test"))
Y_train, Y_test = (np.array(ds[split]["label"], dtype=np.uint8) for split in ("train", "test"))
return X_train, Y_train, X_test, Y_test
def train_model(tm, xtrain, ytrain, xtest, ytest, epochs: int, clause_drop_p: float, model_name: str):
for epoch in range(epochs):
with (fit_timer := Timer()):
loss = tm.fit(xtrain, ytrain, clause_drop_p=clause_drop_p)
with (test_timer := Timer()):
test_pred, _ = tm.predict(xtest)
with (train_timer := Timer()):
train_pred, _ = tm.predict(xtrain)
train_log = {
"Acc": f"{(train_pred == ytrain).mean() * 100:.4f}%",
"Eval Time": f"{train_timer.elapsed:.2f}s",
"Fit Time": f"{fit_timer.elapsed:.2f}s",
}
test_log = {"Acc": f"{(test_pred == ytest).mean() * 100:.4f}%", "Eval Time": f"{test_timer.elapsed:.2f}s"}
if loss is not None: # Only loss-guided TM returns loss
train_log["Loss"] = f"{loss:.4f}"
print_table(
f"{model_name} epoch {epoch + 1}/{epochs}",
{
"Train": train_log,
"Test": test_log,
},
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="TM Kuzushiji-MNIST dataset")
# Subcommand to select discrete or guided
models = parser.add_subparsers(dest="model", required=True)
# Common args
common = argparse.ArgumentParser(add_help=False)
common.add_argument("--epochs", type=int, default=10)
common.add_argument("--clause_drop_p", type=float, default=0.0)
common.add_argument("--n_clauses", type=int, default=500)
common.add_argument("--s", type=float, default=10.0)
common.add_argument("--coalesced", type=lambda v: bool(int(v)), default=True, metavar="0|1")
common.add_argument(
"--patch",
dest="patch_dim",
type=lambda v: None if v.lower() == "none" else tuple(map(int, v.split(","))),
default="10,10",
help="H,W or none",
)
common.add_argument("--seed", type=int, default=10)
common.add_argument("--device", type=str, default="cpu:1", help="cpu:N_THREADS or cuda:GPU_ID")
# Discrete args
discrete = models.add_parser("discrete", parents=[common])
discrete.add_argument("--T", type=float, default=1000)
# Guided args
guided = models.add_parser("guided", parents=[common])
guided.add_argument("--lr", type=float, default=1.0)
guided.add_argument("--lambda", dest="lambda_", type=float, default=1.0)
# Parse args
args = parser.parse_args()
# Args to dict
params = vars(args)
model_name, epochs, clause_drop_p = params.pop("model"), params.pop("epochs"), params.pop("clause_drop_p")
# Model Initialization
TM = DiscreteTM if model_name == "discrete" else GuidedTM
tm = TM(
**params,
dim=(28, 28, 1),
n_classes=10,
)
# Training
train_model(tm, *load_kmnist(), epochs=epochs, clause_drop_p=clause_drop_p, model_name=model_name)
image/cifar10.py#
Train discrete or guided TM on CIFAR-10 dataset.
"""Train discrete or guided TM on CIFAR-10 dataset.
Images are RGB(or gray when --gray).
Binarization per channel:
adaptive - adaptive Gaussian thresholding (window 11, C 2), 1 bit per channel (default).
thermometer - pixels 0-255 -> levels 0-L (--levels, default 16), thermometer encoded by the TM.
Usage:
python cifar10.py {discrete,guided} [options]
Examples:
python cifar10.py discrete --device cuda:0 # non-coalesced, 6000 clauses per class
python cifar10.py discrete --binarization thermometer --device cuda:0
python cifar10.py discrete --coalesced 1 --n_clauses 20000 --device cuda:0
python cifar10.py guided --lr 0.5 --device cuda:0 --epochs 20
Options:
- python cifar10.py <discrete/guided> --help shows all options.
- Common: --epochs --clause_drop_p --gray --binarization adaptive|thermometer --levels --n_clauses --s --coalesced 0|1 --patch H,W|none --seed --device cpu:N|cuda:N
- discrete: --T
- guided: --lr --lambda
"""
import argparse
import numpy as np
from datasets import load_dataset
from scipy.ndimage import gaussian_filter
from ytm.tm import MultiClassTM as DiscreteTM
from ytm.guided import MultiClassTM as GuidedTM
from ytm.utils import Timer, print_table
def adaptive_threshold(X):
# Same as cv2.adaptiveThreshold(ADAPTIVE_THRESH_GAUSSIAN_C, blockSize=11, C=2) per channel:
# pixel is 1 if above the Gaussian-weighted mean of its 11x11 neighbourhood minus 2.
X = X.astype(np.float32)
local_mean = gaussian_filter(X, sigma=(0, 2.0, 2.0, 0), truncate=2.5, mode="nearest")
return (X > local_mean - 2).astype(np.uint8)
def thermometer(X, levels: int):
# Pixels 0-255 -> levels 0-levels, thermometer encoded by the TM.
return (levels * X.astype(np.float32) / 255.0).astype(np.uint8)
def to_gray(X):
# ITU-R BT.601 luma, (N, H, W, 3) -> (N, H, W, 1).
return np.rint(X @ np.array([0.299, 0.587, 0.114], dtype=np.float32))[..., None].astype(np.uint8)
def load_cifar10(binarization: str, levels: int, gray: bool):
ds = load_dataset("uoft-cs/cifar10")
encode = adaptive_threshold if binarization == "adaptive" else lambda X: thermometer(X, levels)
X_train, X_test = (np.array(ds[split]["img"]) for split in ("train", "test"))
if gray:
X_train, X_test = to_gray(X_train), to_gray(X_test)
X_train, X_test = encode(X_train), encode(X_test)
Y_train, Y_test = (np.array(ds[split]["label"], dtype=np.uint8) for split in ("train", "test"))
return X_train, Y_train, X_test, Y_test
def train_model(tm, xtrain, ytrain, xtest, ytest, epochs: int, clause_drop_p: float, model_name: str):
for epoch in range(epochs):
with (fit_timer := Timer()):
loss = tm.fit(xtrain, ytrain, clause_drop_p=clause_drop_p)
with (test_timer := Timer()):
test_pred, _ = tm.predict(xtest)
with (train_timer := Timer()):
train_pred, _ = tm.predict(xtrain)
train_log = {
"Acc": f"{(train_pred == ytrain).mean() * 100:.4f}%",
"Eval Time": f"{train_timer.elapsed:.2f}s",
"Fit Time": f"{fit_timer.elapsed:.2f}s",
}
test_log = {"Acc": f"{(test_pred == ytest).mean() * 100:.4f}%", "Eval Time": f"{test_timer.elapsed:.2f}s"}
if loss is not None: # Only loss-guided TM returns loss
train_log["Loss"] = f"{loss:.4f}"
print_table(
f"{model_name} epoch {epoch + 1}/{epochs}",
{
"Train": train_log,
"Test": test_log,
},
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="TM CIFAR-10 dataset")
# Subcommand to select discrete or guided
models = parser.add_subparsers(dest="model", required=True)
# Common args
common = argparse.ArgumentParser(add_help=False)
common.add_argument("--epochs", type=int, default=10)
common.add_argument("--binarization", choices=["adaptive", "thermometer"], default="adaptive")
common.add_argument("--levels", type=int, default=16, help="thermometer levels")
common.add_argument("--gray", action="store_true", help="use grayscale images instead of RGB")
common.add_argument("--n_clauses", type=int, default=6000)
common.add_argument("--s", type=float, default=10.0)
common.add_argument("--coalesced", type=lambda v: bool(int(v)), default=False, metavar="0|1")
common.add_argument(
"--patch",
dest="patch_dim",
type=lambda v: None if v.lower() == "none" else tuple(map(int, v.split(","))),
default="8,8",
help="H,W or none",
)
common.add_argument("--clause_drop_p", type=float, default=0.5)
common.add_argument("--seed", type=int, default=10)
common.add_argument("--device", type=str, default="cpu:1", help="cpu:N_THREADS or cuda:GPU_ID")
# Discrete args
discrete = models.add_parser("discrete", parents=[common])
discrete.add_argument("--T", type=float, default=48000)
# Guided args
guided = models.add_parser("guided", parents=[common])
guided.add_argument("--lr", type=float, default=1.0)
guided.add_argument("--lambda", dest="lambda_", type=float, default=1.0)
# Parse args
args = parser.parse_args()
# Args to dict
params = vars(args)
model_name, epochs, clause_drop_p = params.pop("model"), params.pop("epochs"), params.pop("clause_drop_p")
binarization, levels, gray = params.pop("binarization"), params.pop("levels"), params.pop("gray")
# Model Initialization
TM = DiscreteTM if model_name == "discrete" else GuidedTM
tm = TM(
**params,
dim=(32, 32, 1 if gray else 3),
n_classes=10,
feat_maxs=1 if binarization == "adaptive" else levels,
)
# Training
train_model(tm, *load_cifar10(binarization, levels, gray), epochs=epochs, clause_drop_p=clause_drop_p, model_name=model_name)
image/celeba.py#
Train discrete or guided TM on CelebA dataset (40 binary face attributes, multilabel).
"""Train discrete or guided TM on CelebA dataset (40 binary face attributes, multilabel).
Images resized to 64x64 RGB (or gray with --gray), pixels encoded as levels 0-L (--levels, default 8, thermometer), random horizontal flip during training.
Usage:
python celeba.py {discrete,guided} [options]
Examples:
python celeba.py discrete --device cuda:0 # 3x3 conv, 80000 clauses
python celeba.py guided --act_loss bce --device cuda:0
python celeba.py guided --n_clauses 20000 --lambda 10 --device cuda:0 --epochs 20
Options:
- python celeba.py <discrete/guided> --help shows all options.
- Common: --epochs --clause_drop_p --gray --levels --n_clauses --s --patch H,W|none --seed --device cpu:N|cuda:N
- discrete: --T --q
- guided: --lr --lambda --act_loss asl|bce
"""
import argparse
from io import BytesIO
import numpy as np
import PIL.Image
from datasets import Image, load_dataset
from sklearn.metrics import average_precision_score, f1_score, roc_auc_score
from ytm.tm import MultiOutputTM as DiscreteTM
from ytm.guided import MultiOutputTM as GuidedTM
from ytm.guided.backends.act_loss import ASL, SigmoidBCE
from ytm.utils import Timer, print_table
ACT_LOSSES = {
"asl": lambda: ASL(gamma_pos=2.0, gamma_neg=4.0, clip=0.01),
"bce": lambda: SigmoidBCE(),
}
def to_gray(X):
# ITU-R BT.601 luma, (N, H, W, 3) -> (N, H, W, 1).
return np.rint(X @ np.array([0.299, 0.587, 0.114], dtype=np.float32))[..., None].astype(np.uint8)
def load_split(split, levels: int, gray: bool):
# Raw JPEG bytes, so the decoder can downscale while decoding (draft).
split = split.select_columns("image").cast_column("image", Image(decode=False))
X = np.empty((len(split), 64, 64, 3), dtype=np.uint8)
for i, row in enumerate(split):
img = PIL.Image.open(BytesIO(row["image"]["bytes"]))
img.draft("RGB", (64, 64))
X[i] = np.asarray(img.convert("RGB").resize((64, 64)))
if gray:
X = to_gray(X)
# Pixels 0-255 -> levels 0-levels, thermometer encoded by the TM.
return (levels * X.astype(np.float32) / 255.0).astype(np.uint8)
def load_celeba(levels: int, gray: bool):
ds = load_dataset("tpremoli/CelebA-attrs")
attrs = [name for name in ds["train"].column_names if name not in ("image", "prompt_string")]
# Attributes are -1/1 -> 0/1.
Y_train, Y_test = ((np.stack([ds[split][a] for a in attrs], axis=1) + 1) // 2 for split in ("train", "test"))
return (
load_split(ds["train"], levels, gray),
Y_train.astype(np.uint8),
load_split(ds["test"], levels, gray),
Y_test.astype(np.uint8),
)
def multilabel_metrics(y, pred, scores):
# Scores rank the samples per label, so raw votes (discrete) and probabilities (guided) both work for the AUCs.
return {
"Macro F1": f"{f1_score(y, pred, average='macro', zero_division=0) * 100:.2f}%",
"Weighted F1": f"{f1_score(y, pred, average='weighted', zero_division=0) * 100:.2f}%",
"ROC AUC": f"{roc_auc_score(y, scores, average='macro') * 100:.2f}%",
"PR AUC": f"{average_precision_score(y, scores, average='macro') * 100:.2f}%",
}
def train_model(tm, xtrain, ytrain, xtest, ytest, epochs: int, clause_drop_p: float, seed: int, model_name: str):
rng = np.random.default_rng(seed)
for epoch in range(epochs):
# Random horizontal flip of half the training images.
flip = rng.random(len(xtrain)) < 0.5
xtrain_epoch = xtrain.copy()
xtrain_epoch[flip] = xtrain_epoch[flip, :, ::-1]
with (fit_timer := Timer()):
loss = tm.fit(xtrain_epoch, ytrain, clause_drop_p=clause_drop_p)
with (test_timer := Timer()):
test_pred, test_scores = tm.predict(xtest)
with (train_timer := Timer()):
train_pred, train_scores = tm.predict(xtrain)
train_log = {
**multilabel_metrics(ytrain, train_pred, tm.to_prob(train_scores)),
"Eval Time": f"{train_timer.elapsed:.2f}s",
"Fit Time": f"{fit_timer.elapsed:.2f}s",
}
test_log = {**multilabel_metrics(ytest, test_pred, tm.to_prob(test_scores)), "Eval Time": f"{test_timer.elapsed:.2f}s"}
if loss is not None: # Only loss-guided TM returns loss
train_log["Loss"] = f"{loss:.4f}"
print_table(
f"{model_name} epoch {epoch + 1}/{epochs}",
{
"Train": train_log,
"Test": test_log,
},
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="TM CelebA dataset")
# Subcommand to select discrete or guided
models = parser.add_subparsers(dest="model", required=True)
# Common args
common = argparse.ArgumentParser(add_help=False)
common.add_argument("--epochs", type=int, default=10)
common.add_argument("--gray", action="store_true", help="use grayscale images instead of RGB")
common.add_argument("--levels", type=int, default=8, help="thermometer levels")
common.add_argument("--clause_drop_p", type=float, default=0.0)
common.add_argument("--n_clauses", type=int, default=80000)
common.add_argument("--s", type=float, default=30.0)
common.add_argument(
"--patch",
dest="patch_dim",
type=lambda v: None if v.lower() == "none" else tuple(map(int, v.split(","))),
default="3,3",
help="H,W or none",
)
common.add_argument("--seed", type=int, default=10)
common.add_argument("--device", type=str, default="cpu:1", help="cpu:N_THREADS or cuda:GPU_ID")
# Discrete args
discrete = models.add_parser("discrete", parents=[common])
discrete.add_argument("--T", type=float, default=80000)
discrete.add_argument("--q", type=float, default=10.0)
# Guided args
guided = models.add_parser("guided", parents=[common])
guided.add_argument("--lr", type=float, default=1.0)
guided.add_argument("--lambda", dest="lambda_", type=float, default=25.0)
guided.add_argument("--act_loss", choices=list(ACT_LOSSES), default="asl")
# Parse args
args = parser.parse_args()
# Args to dict
params = vars(args)
model_name, epochs, clause_drop_p = params.pop("model"), params.pop("epochs"), params.pop("clause_drop_p")
levels, gray = params.pop("levels"), params.pop("gray")
if "act_loss" in params:
params["act_loss"] = ACT_LOSSES[params["act_loss"]]()
# Model Initialization
TM = DiscreteTM if model_name == "discrete" else GuidedTM
tm = TM(
**params,
dim=(64, 64, 1 if gray else 3),
n_classes=40,
feat_maxs=levels,
)
# Training
train_model(tm, *load_celeba(levels, gray), epochs=epochs, clause_drop_p=clause_drop_p, seed=params["seed"], model_name=model_name)
NLP#
nlp/imdb.py#
Train discrete or guided TM on IMDB movie reviews (binary sentiment).
"""Train discrete or guided TM on IMDB movie reviews (binary sentiment).
Text -> binary bag-of-words (unigram + bigram, English stopwords removed), top --max_features terms by chi2.
Usage:
python imdb.py {discrete,guided} [options]
Examples:
python imdb.py discrete --device cuda:0
python imdb.py guided --lr 0.005 --device cuda:0 --epochs 20
python imdb.py discrete --max_features 2000 --n_clauses 1000 --T 2000
Options:
- python imdb.py <discrete/guided> --help shows all options.
- Common: --epochs --clause_drop_p --max_features --n_clauses --s --max_includes --seed --device cpu:N|cuda:N
- discrete: --T
- guided: --lr --lambda
"""
import argparse
import numpy as np
from datasets import load_dataset
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_selection import SelectKBest, chi2
from ytm.tm import BinaryTM as DiscreteTM
from ytm.guided import BinaryTM as GuidedTM
from ytm.utils import Timer, print_table
def load_imdb(max_features: int):
ds = load_dataset("stanfordnlp/imdb")
Y_train, Y_test = (np.array(ds[split]["label"], dtype=np.uint8) for split in ("train", "test"))
vectorizer = CountVectorizer(stop_words="english", ngram_range=(1, 2), binary=True)
selector = SelectKBest(chi2, k=max_features)
X_train = selector.fit_transform(vectorizer.fit_transform(ds["train"]["text"]), Y_train)
X_test = selector.transform(vectorizer.transform(ds["test"]["text"]))
return X_train.toarray().astype(np.uint8), Y_train, X_test.toarray().astype(np.uint8), Y_test
def train_model(tm, xtrain, ytrain, xtest, ytest, epochs: int, clause_drop_p: float, model_name: str):
for epoch in range(epochs):
with (fit_timer := Timer()):
loss = tm.fit(xtrain, ytrain, clause_drop_p=clause_drop_p)
with (test_timer := Timer()):
test_pred, _ = tm.predict(xtest)
with (train_timer := Timer()):
train_pred, _ = tm.predict(xtrain)
train_log = {
"Acc": f"{(train_pred.reshape(-1) == ytrain).mean() * 100:.4f}%",
"Eval Time": f"{train_timer.elapsed:.2f}s",
"Fit Time": f"{fit_timer.elapsed:.2f}s",
}
test_log = {"Acc": f"{(test_pred.reshape(-1) == ytest).mean() * 100:.4f}%", "Eval Time": f"{test_timer.elapsed:.2f}s"}
if loss is not None: # Only loss-guided TM returns loss
train_log["Loss"] = f"{loss:.4f}"
print_table(
f"{model_name} epoch {epoch + 1}/{epochs}",
{
"Train": train_log,
"Test": test_log,
},
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="TM IMDB dataset")
# Subcommand to select discrete or guided
models = parser.add_subparsers(dest="model", required=True)
# Common args
common = argparse.ArgumentParser(add_help=False)
common.add_argument("--epochs", type=int, default=10)
common.add_argument("--clause_drop_p", type=float, default=0.0)
common.add_argument("--max_features", type=int, default=12000, help="terms kept by chi2")
common.add_argument("--n_clauses", type=int, default=10000)
common.add_argument("--s", type=float, default=1.1)
common.add_argument("--max_includes", type=int, default=None, help="max literals per clause (default: no limit)")
common.add_argument("--seed", type=int, default=10)
common.add_argument("--device", type=str, default="cpu:1", help="cpu:N_THREADS or cuda:GPU_ID")
# Discrete args
discrete = models.add_parser("discrete", parents=[common])
discrete.add_argument("--T", type=float, default=10000)
# Guided args
guided = models.add_parser("guided", parents=[common])
guided.add_argument("--lr", type=float, default=0.01)
guided.add_argument("--lambda", dest="lambda_", type=float, default=1.0)
# Parse args
args = parser.parse_args()
# Args to dict
params = vars(args)
model_name, epochs, clause_drop_p = params.pop("model"), params.pop("epochs"), params.pop("clause_drop_p")
max_features = params.pop("max_features")
# Model Initialization
TM = DiscreteTM if model_name == "discrete" else GuidedTM
tm = TM(
**params,
dim=(max_features, 1, 1),
)
# Training
train_model(tm, *load_imdb(max_features), epochs=epochs, clause_drop_p=clause_drop_p, model_name=model_name)
nlp/ag_news.py#
Train discrete or guided TM on AG News (4 news topics).
"""Train discrete or guided TM on AG News (4 news topics).
Text -> lemmatized binary bag-of-words (unigrams, English stopwords removed), top --max_features terms by chi2.
Usage:
python ag_news.py {discrete,guided} [options]
Examples:
python ag_news.py discrete --device cuda:0
python ag_news.py guided --act_loss sce --lr 0.005 --device cuda:0
python ag_news.py guided --lr 0.02 --device cuda:0 --epochs 20
Options:
- python ag_news.py <discrete/guided> --help shows all options.
- Common: --epochs --clause_drop_p --max_features --n_clauses --s --max_includes --seed --device cpu:N|cuda:N
- discrete: --T
- guided: --lr --lambda --act_loss ce|sce
"""
import argparse
import nltk
import numpy as np
from datasets import load_dataset
from nltk.stem import WordNetLemmatizer
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_selection import SelectKBest, chi2
from ytm.tm import MultiClassTM as DiscreteTM
from ytm.guided import MultiClassTM as GuidedTM
from ytm.guided.backends.act_loss import SCE, SoftmaxCE
from ytm.utils import Timer, print_table
ACT_LOSSES = {
"ce": lambda: SoftmaxCE(),
"sce": lambda: SCE(alpha=1.0, beta=1.0),
}
def load_ag_news(max_features: int):
ds = load_dataset("fancyzhx/ag_news")
Y_train, Y_test = (np.array(ds[split]["label"], dtype=np.uint8) for split in ("train", "test"))
nltk.download("wordnet", quiet=True)
lemmatizer = WordNetLemmatizer()
tokenize = CountVectorizer(stop_words="english").build_analyzer()
vectorizer = CountVectorizer(analyzer=lambda doc: [lemmatizer.lemmatize(t) for t in tokenize(doc)], binary=True)
selector = SelectKBest(chi2, k=max_features)
X_train = selector.fit_transform(vectorizer.fit_transform(ds["train"]["text"]), Y_train)
X_test = selector.transform(vectorizer.transform(ds["test"]["text"]))
return X_train.toarray().astype(np.uint8), Y_train, X_test.toarray().astype(np.uint8), Y_test
def train_model(tm, xtrain, ytrain, xtest, ytest, epochs: int, clause_drop_p: float, model_name: str):
for epoch in range(epochs):
with (fit_timer := Timer()):
loss = tm.fit(xtrain, ytrain, clause_drop_p=clause_drop_p)
with (test_timer := Timer()):
test_pred, _ = tm.predict(xtest)
with (train_timer := Timer()):
train_pred, _ = tm.predict(xtrain)
train_log = {
"Acc": f"{(train_pred == ytrain).mean() * 100:.4f}%",
"Eval Time": f"{train_timer.elapsed:.2f}s",
"Fit Time": f"{fit_timer.elapsed:.2f}s",
}
test_log = {"Acc": f"{(test_pred == ytest).mean() * 100:.4f}%", "Eval Time": f"{test_timer.elapsed:.2f}s"}
if loss is not None: # Only loss-guided TM returns loss
train_log["Loss"] = f"{loss:.4f}"
print_table(
f"{model_name} epoch {epoch + 1}/{epochs}",
{
"Train": train_log,
"Test": test_log,
},
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="TM AG News dataset")
# Subcommand to select discrete or guided
models = parser.add_subparsers(dest="model", required=True)
# Common args
common = argparse.ArgumentParser(add_help=False)
common.add_argument("--epochs", type=int, default=10)
common.add_argument("--clause_drop_p", type=float, default=0.0)
common.add_argument("--max_features", type=int, default=2000, help="terms kept by chi2")
common.add_argument("--n_clauses", type=int, default=8000)
common.add_argument("--s", type=float, default=2.0)
common.add_argument("--max_includes", type=int, default=None, help="max literals per clause (default: no limit)")
common.add_argument("--seed", type=int, default=10)
common.add_argument("--device", type=str, default="cpu:1", help="cpu:N_THREADS or cuda:GPU_ID")
# Discrete args
discrete = models.add_parser("discrete", parents=[common])
discrete.add_argument("--T", type=float, default=40000)
# Guided args
guided = models.add_parser("guided", parents=[common])
guided.add_argument("--lr", type=float, default=0.01)
guided.add_argument("--lambda", dest="lambda_", type=float, default=1.0)
guided.add_argument("--act_loss", choices=list(ACT_LOSSES), default="ce")
# Parse args
args = parser.parse_args()
# Args to dict
params = vars(args)
model_name, epochs, clause_drop_p = params.pop("model"), params.pop("epochs"), params.pop("clause_drop_p")
max_features = params.pop("max_features")
if "act_loss" in params:
params["act_loss"] = ACT_LOSSES[params["act_loss"]]()
# Model Initialization
TM = DiscreteTM if model_name == "discrete" else GuidedTM
tm = TM(
**params,
dim=(max_features, 1, 1),
n_classes=4,
)
# Training
train_model(tm, *load_ag_news(max_features), epochs=epochs, clause_drop_p=clause_drop_p, model_name=model_name)
nlp/go_emotions.py#
Train discrete or guided TM on GoEmotions (28 emotion labels, multilabel).
"""Train discrete or guided TM on GoEmotions (28 emotion labels, multilabel).
Text -> tweet tokenizer (case kept, repeated chars collapsed), URLs replaced by a placeholder,
English stopwords removed, lemmatized, binary bag-of-words (unigram + bigram), top --max_features terms by frequency.
Usage:
python go_emotions.py {discrete,guided} [options]
Examples:
python go_emotions.py discrete --device cuda:0
python go_emotions.py guided --device cuda:0 # ASL, lr 0.05, lambda 6
python go_emotions.py guided --act_loss bce --lr 0.5 --lambda 4 --device cuda:0
Options:
- python go_emotions.py <discrete/guided> --help shows all options.
- Common: --epochs --clause_drop_p --max_features --n_clauses --s --max_includes --seed --device cpu:N|cuda:N
- discrete: --T --q
- guided: --lr --lambda --act_loss asl|bce
"""
import argparse
import re
import nltk
import numpy as np
from datasets import load_dataset
from nltk.corpus import stopwords
from nltk.stem import WordNetLemmatizer
from nltk.tokenize import TweetTokenizer
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.metrics import average_precision_score, f1_score, roc_auc_score
from ytm.tm import MultiOutputTM as DiscreteTM
from ytm.guided import MultiOutputTM as GuidedTM
from ytm.guided.backends.act_loss import ASL, SigmoidBCE
from ytm.utils import Timer, print_table
ACT_LOSSES = {
"asl": lambda: ASL(gamma_pos=0.0, gamma_neg=4.0, clip=0.01),
"bce": lambda: SigmoidBCE(),
}
N_LABELS = 28
def load_go_emotions(max_features: int):
ds = load_dataset("google-research-datasets/go_emotions", "simplified")
# Label index lists -> multi-hot (samples, 28).
Y_train, Y_test = (np.zeros((len(ds[split]), N_LABELS), dtype=np.uint8) for split in ("train", "test"))
for split, Y in (("train", Y_train), ("test", Y_test)):
for i, labels in enumerate(ds[split]["labels"]):
Y[i, labels] = 1
nltk.download("stopwords", quiet=True)
nltk.download("wordnet", quiet=True)
stop = set(stopwords.words("english"))
lemmatizer = WordNetLemmatizer()
tokenizer = TweetTokenizer(preserve_case=True, reduce_len=True)
url = re.compile(r"https?://\S+|www\.\S+")
def analyze(doc):
tokens = [lemmatizer.lemmatize(t) for t in tokenizer.tokenize(url.sub("URL", doc)) if t.lower() not in stop]
return tokens + [f"{a} {b}" for a, b in zip(tokens, tokens[1:])]
vectorizer = CountVectorizer(analyzer=analyze, binary=True, max_features=max_features)
X_train = vectorizer.fit_transform(ds["train"]["text"])
X_test = vectorizer.transform(ds["test"]["text"])
return X_train.toarray().astype(np.uint8), Y_train, X_test.toarray().astype(np.uint8), Y_test
def multilabel_metrics(y, pred, scores):
# Scores rank the samples per label, so raw votes (discrete) and probabilities (guided) both work for the AUCs.
return {
"Macro F1": f"{f1_score(y, pred, average='macro', zero_division=0) * 100:.2f}%",
"Weighted F1": f"{f1_score(y, pred, average='weighted', zero_division=0) * 100:.2f}%",
"ROC AUC": f"{roc_auc_score(y, scores, average='macro') * 100:.2f}%",
"PR AUC": f"{average_precision_score(y, scores, average='macro') * 100:.2f}%",
}
def train_model(tm, xtrain, ytrain, xtest, ytest, epochs: int, clause_drop_p: float, model_name: str):
for epoch in range(epochs):
with (fit_timer := Timer()):
loss = tm.fit(xtrain, ytrain, clause_drop_p=clause_drop_p)
with (test_timer := Timer()):
test_pred, test_scores = tm.predict(xtest)
with (train_timer := Timer()):
train_pred, train_scores = tm.predict(xtrain)
train_log = {
**multilabel_metrics(ytrain, train_pred, tm.to_prob(train_scores)),
"Eval Time": f"{train_timer.elapsed:.2f}s",
"Fit Time": f"{fit_timer.elapsed:.2f}s",
}
test_log = {**multilabel_metrics(ytest, test_pred, tm.to_prob(test_scores)), "Eval Time": f"{test_timer.elapsed:.2f}s"}
if loss is not None: # Only loss-guided TM returns loss
train_log["Loss"] = f"{loss:.4f}"
print_table(
f"{model_name} epoch {epoch + 1}/{epochs}",
{
"Train": train_log,
"Test": test_log,
},
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="TM GoEmotions dataset")
# Subcommand to select discrete or guided
models = parser.add_subparsers(dest="model", required=True)
# Common args
common = argparse.ArgumentParser(add_help=False)
common.add_argument("--epochs", type=int, default=10)
common.add_argument("--clause_drop_p", type=float, default=0.0)
common.add_argument("--max_features", type=int, default=10000, help="most frequent terms kept")
common.add_argument("--n_clauses", type=int, default=22400)
common.add_argument("--s", type=float, default=2.0)
common.add_argument("--max_includes", type=int, default=None, help="max literals per clause (default: no limit)")
common.add_argument("--seed", type=int, default=10)
common.add_argument("--device", type=str, default="cpu:1", help="cpu:N_THREADS or cuda:GPU_ID")
# Discrete args
discrete = models.add_parser("discrete", parents=[common])
discrete.add_argument("--T", type=float, default=25000)
discrete.add_argument("--q", type=float, default=1.5)
# Guided args
guided = models.add_parser("guided", parents=[common])
guided.add_argument("--lr", type=float, default=0.05)
guided.add_argument("--lambda", dest="lambda_", type=float, default=6.0)
guided.add_argument("--act_loss", choices=list(ACT_LOSSES), default="asl")
# Parse args
args = parser.parse_args()
# Args to dict
params = vars(args)
model_name, epochs, clause_drop_p = params.pop("model"), params.pop("epochs"), params.pop("clause_drop_p")
max_features = params.pop("max_features")
if "act_loss" in params:
params["act_loss"] = ACT_LOSSES[params["act_loss"]]()
# Model Initialization
TM = DiscreteTM if model_name == "discrete" else GuidedTM
tm = TM(
**params,
dim=(max_features, 1, 1),
n_classes=N_LABELS,
)
# Training
train_model(tm, *load_go_emotions(max_features), epochs=epochs, clause_drop_p=clause_drop_p, model_name=model_name)
Regression#
reg/friedman1.py#
Train discrete or guided regression TM on Friedman #1 (synthetic, sklearn make_friedman1).
"""Train discrete or guided regression TM on Friedman #1 (synthetic, sklearn make_friedman1).
Usage:
python friedman1.py {discrete,guided} [options]
Examples:
python friedman1.py discrete
python friedman1.py discrete --n_features 20 --N 10000 --device cuda:0
python friedman1.py guided --n_clauses 1000
Options:
- python friedman1.py <discrete/guided> --help shows all options.
- Common: --N --n_features --noise --epochs --bins --strategy uniform|quantile --n_clauses --s --max_includes --boost_tp_inc 0|1 --boost_tp_dec 0|1 --seed --device cpu:N|cuda:N
- discrete: --T
- guided: --lr --lambda --act_loss mse|mae|huber
"""
import argparse
import numpy as np
from sklearn.datasets import make_friedman1
from sklearn.metrics import mean_absolute_error, r2_score, root_mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import KBinsDiscretizer
from ytm.tm import RegressionTM as DiscreteTM
from ytm.guided import RegressionTM as GuidedTM
from ytm.guided.backends.act_loss import MAE, MSE, Huber
from ytm.utils import Timer, print_table
ACT_LOSSES = {"mse": MSE, "mae": MAE, "huber": Huber}
def load(N: int, n_features: int, noise: float, bins: int, strategy: str, seed: int):
X, Y = make_friedman1(n_samples=N, n_features=n_features, noise=noise, random_state=seed)
X_train, X_test, Y_train, Y_test = train_test_split(X, Y.astype(np.float32), test_size=0.3, random_state=seed)
d = KBinsDiscretizer(n_bins=bins, encode="ordinal", strategy=strategy)
X_train, X_test = d.fit_transform(X_train).astype(np.int32), d.transform(X_test).astype(np.int32)
return X_train, Y_train, X_test, Y_test, d.n_bins_ - 1
def regression_metrics(y, pred):
return {
"RMSE": f"{root_mean_squared_error(y, pred):.3f}",
"MAE": f"{mean_absolute_error(y, pred):.3f}",
"R2": f"{r2_score(y, pred):.3f}",
}
def train_model(tm, xtrain, ytrain, xtest, ytest, epochs: int, model_name: str):
# Standardize targets with the training mean/std, metrics are reported in original units.
mu, sd = ytrain.mean(), ytrain.std()
ytrain_std = (ytrain - mu) / sd
for epoch in range(epochs):
with (fit_timer := Timer()):
loss = tm.fit(xtrain, ytrain_std)
with (test_timer := Timer()):
test_pred = tm.predict(xtest)[0] * sd + mu
with (train_timer := Timer()):
train_pred = tm.predict(xtrain)[0] * sd + mu
train_log = {
**regression_metrics(ytrain, train_pred),
"Eval Time": f"{train_timer.elapsed:.2f}s",
"Fit Time": f"{fit_timer.elapsed:.2f}s",
}
test_log = {**regression_metrics(ytest, test_pred), "Eval Time": f"{test_timer.elapsed:.2f}s"}
if loss is not None: # Only loss-guided TM returns loss
train_log["Loss"] = f"{loss:.4f}"
print_table(
f"{model_name} epoch {epoch + 1}/{epochs}",
{
"Train": train_log,
"Test": test_log,
},
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="TM Friedman #1 regression")
# Subcommand to select discrete or guided
models = parser.add_subparsers(dest="model", required=True)
# Common args
common = argparse.ArgumentParser(add_help=False)
common.add_argument("--N", type=int, default=5000, help="dataset size (70/30 train/test)")
common.add_argument("--n_features", type=int, default=10, help="total features, first 5 are relevant")
common.add_argument("--noise", type=float, default=1.0, help="std of gaussian noise on y")
common.add_argument("--epochs", type=int, default=100)
common.add_argument("--bins", type=int, default=32, help="discretization levels per feature")
common.add_argument("--strategy", choices=["uniform", "quantile"], default="quantile")
common.add_argument("--n_clauses", type=int, default=100)
common.add_argument("--s", type=float, default=5.0)
common.add_argument("--max_includes", type=int, default=8, help="max literals per clause")
common.add_argument("--boost_tp_inc", type=lambda v: bool(int(v)), default=False, metavar="0|1")
common.add_argument("--boost_tp_dec", type=lambda v: bool(int(v)), default=False, metavar="0|1")
common.add_argument("--seed", type=int, default=10)
common.add_argument("--device", type=str, default="cpu:1", help="cpu:N_THREADS or cuda:GPU_ID")
# Discrete args
discrete = models.add_parser("discrete", parents=[common])
discrete.add_argument("--T", type=float, default=5000)
# Guided args
guided = models.add_parser("guided", parents=[common])
guided.add_argument("--lr", type=float, default=0.01)
guided.add_argument("--lambda", dest="lambda_", type=float, default=30.0)
guided.add_argument("--act_loss", choices=list(ACT_LOSSES), default="mse")
# Parse args
args = parser.parse_args()
# Args to dict
params = vars(args)
model_name, epochs = params.pop("model"), params.pop("epochs")
N, n_features, noise = params.pop("N"), params.pop("n_features"), params.pop("noise")
bins, strategy = params.pop("bins"), params.pop("strategy")
if "act_loss" in params:
params["act_loss"] = ACT_LOSSES[params["act_loss"]]()
X_train, Y_train, X_test, Y_test, feat_maxs = load(N, n_features, noise, bins, strategy, params["seed"])
# Model Initialization.
if model_name == "discrete":
y_std = (Y_train - Y_train.mean()) / Y_train.std()
tm = DiscreteTM(**params, dim=(n_features, 1, 1), y_range=(float(y_std.min()), float(y_std.max())), feat_maxs=feat_maxs)
else:
tm = GuidedTM(**params, dim=(n_features, 1, 1), feat_maxs=feat_maxs)
# Training
train_model(tm, X_train, Y_train, X_test, Y_test, epochs=epochs, model_name=model_name)
reg/diabetes.py#
Train discrete or guided regression TM on Diabetes (442 patients, disease progression after one year).
"""Train discrete or guided regression TM on Diabetes (442 patients, disease progression after one year).
Usage:
python diabetes.py {discrete,guided} [options]
Examples:
python diabetes.py discrete
python diabetes.py discrete --bins 4 --strategy uniform
python diabetes.py guided --act_loss huber
Options:
- python diabetes.py <discrete/guided> --help shows all options.
- Common: --epochs --bins --strategy uniform|quantile --n_clauses --s --max_includes --boost_tp_inc 0|1 --boost_tp_dec 0|1 --seed --device cpu:N|cuda:N
- discrete: --T
- guided: --lr --lambda --act_loss mse|mae|huber
"""
import argparse
import numpy as np
from sklearn.datasets import load_diabetes
from sklearn.metrics import mean_absolute_error, r2_score, root_mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import KBinsDiscretizer
from ytm.tm import RegressionTM as DiscreteTM
from ytm.guided import RegressionTM as GuidedTM
from ytm.guided.backends.act_loss import MAE, MSE, Huber
from ytm.utils import Timer, print_table
ACT_LOSSES = {"mse": MSE, "mae": MAE, "huber": Huber}
def load(bins: int, strategy: str, seed: int):
X, Y = load_diabetes(return_X_y=True, scaled=False)
X_train, X_test, Y_train, Y_test = train_test_split(X, Y.astype(np.float32), test_size=0.3, random_state=seed)
# Continuous features: per-feature bins. Sex: 1/2 -> 0/1.
cont = [i for i in range(X.shape[1]) if i != 1]
d = KBinsDiscretizer(n_bins=bins, encode="ordinal", strategy=strategy)
d.fit(X_train[:, cont])
def encode(X):
out = np.empty(X.shape, dtype=np.int32)
out[:, cont] = d.transform(X[:, cont])
out[:, 1] = X[:, 1] - 1 # feat: sex = 0|1
return out
# Highest level per feature: bins-1 for the discretized ones, 1 for sex.
feat_maxs = np.empty(X.shape[1], dtype=np.int32)
feat_maxs[cont] = d.n_bins_ - 1
feat_maxs[1] = 1
return encode(X_train), Y_train, encode(X_test), Y_test, feat_maxs
def regression_metrics(y, pred):
return {
"RMSE": f"{root_mean_squared_error(y, pred):.3f}",
"MAE": f"{mean_absolute_error(y, pred):.3f}",
"R2": f"{r2_score(y, pred):.3f}",
}
def train_model(tm, xtrain, ytrain, xtest, ytest, epochs: int, model_name: str):
# Standardize targets with the training mean/std, metrics are reported in original units.
mu, sd = ytrain.mean(), ytrain.std()
ytrain_std = (ytrain - mu) / sd
for epoch in range(epochs):
with (fit_timer := Timer()):
loss = tm.fit(xtrain, ytrain_std)
with (test_timer := Timer()):
test_pred = tm.predict(xtest)[0] * sd + mu
with (train_timer := Timer()):
train_pred = tm.predict(xtrain)[0] * sd + mu
train_log = {
**regression_metrics(ytrain, train_pred),
"Eval Time": f"{train_timer.elapsed:.2f}s",
"Fit Time": f"{fit_timer.elapsed:.2f}s",
}
test_log = {**regression_metrics(ytest, test_pred), "Eval Time": f"{test_timer.elapsed:.2f}s"}
if loss is not None: # Only loss-guided TM returns loss
train_log["Loss"] = f"{loss:.4f}"
print_table(
f"{model_name} epoch {epoch + 1}/{epochs}",
{
"Train": train_log,
"Test": test_log,
},
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="TM Diabetes regression")
# Subcommand to select discrete or guided
models = parser.add_subparsers(dest="model", required=True)
# Common args
common = argparse.ArgumentParser(add_help=False)
common.add_argument("--epochs", type=int, default=50)
common.add_argument("--bins", type=int, default=6, help="discretization levels per continuous feature")
common.add_argument("--strategy", choices=["uniform", "quantile"], default="quantile")
common.add_argument("--n_clauses", type=int, default=100)
common.add_argument("--s", type=float, default=2.0)
common.add_argument("--max_includes", type=int, default=None, help="max literals per clause (default: no limit)")
common.add_argument("--boost_tp_inc", type=lambda v: bool(int(v)), default=False, metavar="0|1")
common.add_argument("--boost_tp_dec", type=lambda v: bool(int(v)), default=False, metavar="0|1")
common.add_argument("--seed", type=int, default=10)
common.add_argument("--device", type=str, default="cpu:1", help="cpu:N_THREADS or cuda:GPU_ID")
# Discrete args
discrete = models.add_parser("discrete", parents=[common])
discrete.add_argument("--T", type=float, default=300)
# Guided args
guided = models.add_parser("guided", parents=[common])
guided.add_argument("--lr", type=float, default=0.03)
guided.add_argument("--lambda", dest="lambda_", type=float, default=1.0)
guided.add_argument("--act_loss", choices=list(ACT_LOSSES), default="mse")
# Parse args
args = parser.parse_args()
# Args to dict
params = vars(args)
model_name, epochs = params.pop("model"), params.pop("epochs")
bins, strategy = params.pop("bins"), params.pop("strategy")
if "act_loss" in params:
params["act_loss"] = ACT_LOSSES[params["act_loss"]]()
X_train, Y_train, X_test, Y_test, feat_maxs = load(bins, strategy, params["seed"])
# Model Initialization.
n_features = X_train.shape[1]
if model_name == "discrete":
y_std = (Y_train - Y_train.mean()) / Y_train.std()
tm = DiscreteTM(**params, dim=(n_features, 1, 1), y_range=(float(y_std.min()), float(y_std.max())), feat_maxs=feat_maxs)
else:
tm = GuidedTM(**params, dim=(n_features, 1, 1), feat_maxs=feat_maxs)
# Training
train_model(tm, X_train, Y_train, X_test, Y_test, epochs=epochs, model_name=model_name)
reg/california_housing.py#
Train discrete or guided regression TM on California Housing (8 numeric features, median house value in $100k).
"""Train discrete or guided regression TM on California Housing (8 numeric features, median house value in $100k).
Data from the Hugging Face mirror gvlassis/california_housing (same 20640 rows as sklearn).
Usage:
python california_housing.py {discrete,guided} [options]
Examples:
python california_housing.py discrete --device cuda:0
python california_housing.py guided --n_clauses 3000 --device cuda:0
python california_housing.py guided --act_loss huber
Options:
- python california_housing.py <discrete/guided> --help shows all options.
- Common: --epochs --bins --strategy uniform|quantile --n_clauses --s --max_includes --boost_tp_inc 0|1 --boost_tp_dec 0|1 --seed --device cpu:N|cuda:N
- discrete: --T
- guided: --lr --lambda --act_loss mse|mae|huber
"""
import argparse
import numpy as np
from datasets import concatenate_datasets, load_dataset
from sklearn.metrics import mean_absolute_error, r2_score, root_mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import KBinsDiscretizer
from ytm.tm import RegressionTM as DiscreteTM
from ytm.guided import RegressionTM as GuidedTM
from ytm.guided.backends.act_loss import MAE, MSE, Huber
from ytm.utils import Timer, print_table
ACT_LOSSES = {"mse": MSE, "mae": MAE, "huber": Huber}
def load(bins: int, strategy: str, seed: int):
ds = load_dataset("gvlassis/california_housing")
ds = concatenate_datasets([ds["train"], ds["validation"], ds["test"]])
X = np.stack([np.array(ds[c]) for c in ds.column_names if c != "MedHouseVal"], axis=1)
Y = np.array(ds["MedHouseVal"], dtype=np.float32)
X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.3, random_state=seed)
d = KBinsDiscretizer(n_bins=bins, encode="ordinal", strategy=strategy)
X_train, X_test = d.fit_transform(X_train).astype(np.int32), d.transform(X_test).astype(np.int32)
return X_train, Y_train, X_test, Y_test, d.n_bins_ - 1
def regression_metrics(y, pred):
return {
"RMSE": f"{root_mean_squared_error(y, pred):.3f}",
"MAE": f"{mean_absolute_error(y, pred):.3f}",
"R2": f"{r2_score(y, pred):.3f}",
}
def train_model(tm, xtrain, ytrain, xtest, ytest, epochs: int, model_name: str):
# Standardize targets with the training mean/std, metrics are reported in original units.
mu, sd = ytrain.mean(), ytrain.std()
ytrain_std = (ytrain - mu) / sd
for epoch in range(epochs):
with (fit_timer := Timer()):
loss = tm.fit(xtrain, ytrain_std)
with (test_timer := Timer()):
test_pred = tm.predict(xtest)[0] * sd + mu
with (train_timer := Timer()):
train_pred = tm.predict(xtrain)[0] * sd + mu
train_log = {
**regression_metrics(ytrain, train_pred),
"Eval Time": f"{train_timer.elapsed:.2f}s",
"Fit Time": f"{fit_timer.elapsed:.2f}s",
}
test_log = {**regression_metrics(ytest, test_pred), "Eval Time": f"{test_timer.elapsed:.2f}s"}
if loss is not None: # Only loss-guided TM returns loss
train_log["Loss"] = f"{loss:.4f}"
print_table(
f"{model_name} epoch {epoch + 1}/{epochs}",
{
"Train": train_log,
"Test": test_log,
},
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="TM California Housing regression")
# Subcommand to select discrete or guided
models = parser.add_subparsers(dest="model", required=True)
# Common args
common = argparse.ArgumentParser(add_help=False)
common.add_argument("--epochs", type=int, default=150)
common.add_argument("--bins", type=int, default=64, help="discretization levels per feature")
common.add_argument("--strategy", choices=["uniform", "quantile"], default="quantile")
common.add_argument("--n_clauses", type=int, default=250)
common.add_argument("--s", type=float, default=20.0)
common.add_argument("--max_includes", type=int, default=None, help="max literals per clause (default: no limit)")
common.add_argument("--boost_tp_inc", type=lambda v: bool(int(v)), default=True, metavar="0|1")
common.add_argument("--boost_tp_dec", type=lambda v: bool(int(v)), default=False, metavar="0|1")
common.add_argument("--seed", type=int, default=10)
common.add_argument("--device", type=str, default="cpu:1", help="cpu:N_THREADS or cuda:GPU_ID")
# Discrete args
discrete = models.add_parser("discrete", parents=[common])
discrete.add_argument("--T", type=float, default=20000)
# Guided args
guided = models.add_parser("guided", parents=[common])
guided.add_argument("--lr", type=float, default=0.03)
guided.add_argument("--lambda", dest="lambda_", type=float, default=30.0)
guided.add_argument("--act_loss", choices=list(ACT_LOSSES), default="mse")
# Parse args
args = parser.parse_args()
# Args to dict
params = vars(args)
model_name, epochs = params.pop("model"), params.pop("epochs")
bins, strategy = params.pop("bins"), params.pop("strategy")
if "act_loss" in params:
params["act_loss"] = ACT_LOSSES[params["act_loss"]]()
X_train, Y_train, X_test, Y_test, feat_maxs = load(bins, strategy, params["seed"])
# Model Initialization.
n_features = X_train.shape[1]
if model_name == "discrete":
y_std = (Y_train - Y_train.mean()) / Y_train.std()
tm = DiscreteTM(**params, dim=(n_features, 1, 1), y_range=(float(y_std.min()), float(y_std.max())), feat_maxs=feat_maxs)
else:
tm = GuidedTM(**params, dim=(n_features, 1, 1), feat_maxs=feat_maxs)
# Training
train_model(tm, X_train, Y_train, X_test, Y_test, epochs=epochs, model_name=model_name)