wehub-resource-sync
144 lines
4.3 KiB
Python
144 lines
4.3 KiB
Python
"""
|
|
Scripts to test cleanlab usage with various ML frameworks:
|
|
pytorch, skorch
|
|
"""
|
|
|
|
import pytest
|
|
import warnings
|
|
|
|
# pytest.mark.filterwarnings is unable to catch filterbuffers library DeprecationWarning
|
|
warnings.filterwarnings(action="ignore", category=DeprecationWarning)
|
|
|
|
import os
|
|
from copy import deepcopy
|
|
import random
|
|
import numpy as np
|
|
|
|
import torch
|
|
import skorch
|
|
import sklearn
|
|
|
|
from cleanlab.classification import CleanLearning
|
|
|
|
# Version check for sklearn compatibility
|
|
uses_sklearn_ge_1_6_0 = tuple(map(int, sklearn.__version__.split(".")[:2])) >= (1, 6)
|
|
|
|
|
|
def dataset_w_errors():
|
|
num_classes = 2
|
|
num_features = 3
|
|
n = 50
|
|
margin = 5
|
|
X = np.vstack(
|
|
[np.random.random((n, num_features)), np.random.random((n, num_features)) + margin]
|
|
)
|
|
X = (X - X.mean(axis=0)) / X.std(axis=0) # normalize columns
|
|
y = np.array([0] * n + [1] * n)
|
|
y_og = np.array(y)
|
|
# Introduce label errors
|
|
error_indices = [n - 3, n - 2, n - 1, n, n + 1, n + 2]
|
|
for idx in error_indices:
|
|
y[idx] = 1 - y[idx] # Flip label
|
|
|
|
if os.name == "nt": # running on Windows
|
|
# numpy converts to int32 instead of int64 on Windows, incompatible with neural nets
|
|
# https://github.com/numpy/numpy/issues/17640
|
|
X = np.float64(X)
|
|
y = np.int64(y)
|
|
y_og = np.int64(y_og)
|
|
|
|
return {
|
|
"X": X,
|
|
"y": y,
|
|
"y_og": y_og,
|
|
"error_indices": error_indices,
|
|
"num_classes": num_classes,
|
|
"num_features": num_features,
|
|
}
|
|
|
|
|
|
def make_rare_label(data):
|
|
"""Makes one label really rare in the dataset."""
|
|
data = deepcopy(data)
|
|
y = data["y"]
|
|
class0_inds = np.where(y == 0)[0]
|
|
if len(class0_inds) < 1:
|
|
raise ValueError("Class 0 too rare already")
|
|
class0_inds_remove = class0_inds[1:]
|
|
if len(class0_inds_remove) > 0:
|
|
y[class0_inds_remove] = 1
|
|
data["y"] = y
|
|
return data
|
|
|
|
|
|
SEED = 1
|
|
np.random.seed(SEED)
|
|
random.seed(SEED)
|
|
torch.manual_seed(SEED)
|
|
torch.backends.cudnn.deterministic = True
|
|
torch.backends.cudnn.benchmark = False
|
|
torch.cuda.manual_seed_all(SEED)
|
|
|
|
DATA = dataset_w_errors()
|
|
DATA_RARE_LABEL = make_rare_label(DATA)
|
|
|
|
|
|
def test_torch(data=DATA, hidden_units=128):
|
|
dataset = torch.utils.data.TensorDataset(
|
|
torch.from_numpy(data["X"]).float(), torch.from_numpy(data["y"])
|
|
)
|
|
|
|
class TorchNetwork(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.layers = torch.nn.Sequential(
|
|
torch.nn.Linear(data["num_features"], hidden_units),
|
|
torch.nn.ReLU(),
|
|
torch.nn.Linear(hidden_units, data["num_classes"]),
|
|
)
|
|
|
|
def forward(self, X):
|
|
return self.layers(X)
|
|
|
|
# Test base model works:
|
|
skorch_config = {"criterion": torch.nn.CrossEntropyLoss, "optimizer": torch.optim.Adam}
|
|
net = skorch.NeuralNet(TorchNetwork, **skorch_config)
|
|
net.fit(dataset, data["y"], epochs=2)
|
|
preds_base = net.predict(dataset)
|
|
|
|
# Test Cleanlearning performs well:
|
|
net = skorch.NeuralNet(TorchNetwork, **skorch_config)
|
|
cl = CleanLearning(net)
|
|
cl.fit(dataset, data["y"], clf_kwargs={"epochs": 30}, clf_final_kwargs={"epochs": 60})
|
|
|
|
preds = cl.predict(dataset).argmax(axis=1)
|
|
err = np.sum(preds != data["y_og"]) / len(data["y_og"])
|
|
issue_indices = list(cl.label_issues_df[cl.label_issues_df["is_label_issue"]].index.values)
|
|
assert issue_indices == data["error_indices"]
|
|
assert err < 1e-3
|
|
|
|
|
|
@pytest.mark.filterwarnings("ignore")
|
|
def test_torch_rarelabel(data=DATA_RARE_LABEL, hidden_units=8):
|
|
dataset = torch.utils.data.TensorDataset(
|
|
torch.from_numpy(data["X"]).float(), torch.from_numpy(data["y"])
|
|
)
|
|
|
|
class TorchNetwork(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.layers = torch.nn.Sequential(
|
|
torch.nn.Linear(data["num_features"], hidden_units),
|
|
torch.nn.ReLU(),
|
|
torch.nn.Linear(hidden_units, data["num_classes"]),
|
|
)
|
|
|
|
def forward(self, X):
|
|
return self.layers(X)
|
|
|
|
# Test Cleanlearning works:
|
|
net = skorch.NeuralNet(TorchNetwork, criterion=torch.nn.CrossEntropyLoss)
|
|
cl = CleanLearning(net)
|
|
cl.fit(dataset, data["y"], clf_kwargs={"epochs": 2})
|
|
pred_probs = cl.predict(dataset)
|