Improving the iris dataset classification

Hi, Thanks for the library. I made a small change in the iris classification demo, using the full dataset (i.e. num_features=4) and then scaling the dataset using sklearn.preprocessing import StandardScaler, the classifier now achieves 100% training accuracy and 98% test accuracy!

Here's the code:

import torch
from matplotlib.colors import ListedColormap
from sklearn import datasets
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from simple_einet.einet import Einet, EinetConfig
from simple_einet.layers.distributions.normal import Normal, RatNormal
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

# Load the Iris dataset
iris = datasets.load_iris()
X = iris.data
scaler = StandardScaler()
X = scaler.fit_transform(X)
# Split the dataset into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(X, iris.target, test_size=0.33, random_state=42)

# Convert the data to PyTorch tensors
X_train = torch.tensor(X_train).float()
y_train = torch.tensor(y_train).long()
X_test = torch.tensor(X_test).float()
y_test = torch.tensor(y_test).long()

# Define model configuration using default values
config = EinetConfig(
    num_features=4,
    depth=1,
    num_sums=5,
    num_leaves=5,
    num_repetitions=5,
    num_classes=3,
    leaf_type=Normal,
    leaf_kwargs={},
    dropout=0.0,
)

# Initialize the model
model = Einet(config)

# Set up the optimizer and loss function
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
cross_entropy = torch.nn.CrossEntropyLoss()

# Define a function to compute accuracy
def accuracy(model, X, y):
    with torch.no_grad():
        outputs = model(X)
        predictions = outputs.argmax(-1)
        correct = (predictions == y).sum()
        total = y.shape[0]
        return correct / total * 100

# Training loop
for epoch in range(50):
    optimizer.zero_grad()
    log_likelihoods = model(X_train)
    loss = cross_entropy(log_likelihoods, y_train)
    loss.backward()
    optimizer.step()

    acc_train = accuracy(model, X_train, y_train)
    acc_test = accuracy(model, X_test, y_test)
    print(f"Epoch: {epoch + 1}, Loss: {loss.item():.2f}, Accuracy Train: {acc_train:.2f} %, Accuracy Test: {acc_test:.2f} %")

braun-steven / simple-einet

Improving the iris dataset classification #6