機械学習においては学習に使用するデータの量だけでなく質を担保することが非常に重要ということで、昔は自分で論文を参考にアノテーションミス候補のデータを絞るロジックを書いていたのですが、今なら何かライブラリがあるだろうと思いcleanlabというライブラリを見つけたので、そのライブラリの使い方を記す簡単なサンプルをcodexに作成してもらったのでメモします。
cleanlabにfind_label_issuesという関数があって、そこにアノテーションや推論によって出力された確率を渡してあげれば良いようです。
結果として"アノテーションミスではないか"と出力されたHTMLは以下の通りで、確かにアノテーションミスが疑われるケースを抽出できている。
cleanlabによってアノテーションミスと疑われた例
detect_mislabeled_samples.py
import numpy as np
import torch
from cleanlab.filter import find_label_issues
from sklearn.model_selection import StratifiedKFold
from torch.utils.data import DataLoader, Subset
from mnist_dataset import MNISTConfig, MNISTDataset
from train import train_one_fold
from visualizer import save_suspicious_html
def predict_probabilities(model, data_loader, device: str) -> np.ndarray:
model.eval()
probabilities = []
with torch.no_grad():
for images, _ in data_loader:
images = images.to(device)
logits = model(images)
batch_probs = torch.softmax(logits, dim=1)
probabilities.append(batch_probs.cpu().numpy())
return np.concatenate(probabilities, axis=0)
def predict_all_data_with_kfold(
dataset,
labels: np.ndarray,
n_splits: int,
epochs: int,
batch_size: int,
learning_rate: float,
device: str,
) -> np.ndarray:
pred_probs = np.zeros((len(dataset), 10), dtype=np.float32)
kfold = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=0)
for fold, (train_indices, valid_indices) in enumerate(
kfold.split(np.zeros(len(labels)), labels),
start=1,
):
print(f"fold={fold}/{n_splits} train={len(train_indices)} valid={len(valid_indices)}")
train_subset = Subset(dataset, train_indices.tolist())
valid_subset = Subset(dataset, valid_indices.tolist())
valid_loader = DataLoader(valid_subset, batch_size=batch_size, shuffle=False)
model = train_one_fold(
train_subset=train_subset,
batch_size=batch_size,
epochs=epochs,
learning_rate=learning_rate,
device=device,
)
pred_probs[valid_indices] = predict_probabilities(model, valid_loader, device)
return pred_probs
def detect_mislabeled_samples(
top_n: int = 100,
n_splits: int = 5,
epochs: int = 10,
learning_rate: float = 1e-3,
) -> dict[str, np.ndarray]:
device = "cuda" if torch.cuda.is_available() else "cpu"
mnist = MNISTDataset(MNISTConfig())
dataset = mnist.train_dataset()
labels = np.array(MNISTDataset.labels(dataset))
pred_probs = predict_all_data_with_kfold(
dataset=dataset,
labels=labels,
n_splits=n_splits,
epochs=epochs,
batch_size=mnist.config.batch_size,
learning_rate=learning_rate,
device=device,
)
ranked_indices = find_label_issues(
labels=labels,
pred_probs=pred_probs,
return_indices_ranked_by="self_confidence",
)
suspicious_indices = ranked_indices[:top_n]
return {
"indices": suspicious_indices,
"labels": labels,
"pred_probs": pred_probs,
}
if __name__ == "__main__":
result = detect_mislabeled_samples()
output_path = save_suspicious_html(
indices=result["indices"],
labels=result["labels"],
pred_probs=result["pred_probs"],
)
print(f"saved visualization: {output_path}")
mnist_dataset.py
from dataclasses import dataclass
from torch.utils.data import DataLoader, Dataset, Subset
from torchvision import datasets, transforms
@dataclass(frozen=True)
class MNISTConfig:
"""Settings shared by training and label-issue detection."""
data_dir: str = "data"
batch_size: int = 128
num_workers: int = 0
train_limit: int | None = None
test_limit: int | None = None
class MNISTDataset:
"""Small wrapper around torchvision's MNIST dataset.
Keeping dataset setup here makes the other scripts focus on their own job:
training, cleanlab scoring, or visualization.
"""
def __init__(self, config: MNISTConfig | None = None):
self.config = config or MNISTConfig()
self.transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,)),
]
)
def train_dataset(self) -> Dataset:
dataset = datasets.MNIST(
root=self.config.data_dir,
train=True,
download=True,
transform=self.transform,
)
return self._limit(dataset, self.config.train_limit)
def test_dataset(self) -> Dataset:
dataset = datasets.MNIST(
root=self.config.data_dir,
train=False,
download=True,
transform=self.transform,
)
return self._limit(dataset, self.config.test_limit)
def train_loader(self, shuffle: bool = True) -> DataLoader:
return DataLoader(
self.train_dataset(),
batch_size=self.config.batch_size,
shuffle=shuffle,
num_workers=self.config.num_workers,
)
def test_loader(self) -> DataLoader:
return DataLoader(
self.test_dataset(),
batch_size=self.config.batch_size,
shuffle=False,
num_workers=self.config.num_workers,
)
@staticmethod
def labels(dataset: Dataset) -> list[int]:
"""Return labels for a Dataset or Subset without loading all images."""
if isinstance(dataset, Subset):
base_targets = dataset.dataset.targets
return [int(base_targets[i]) for i in dataset.indices]
return [int(label) for label in dataset.targets]
@staticmethod
def _limit(dataset: Dataset, limit: int | None) -> Dataset:
if limit is None:
return dataset
return Subset(dataset, range(min(limit, len(dataset))))
simple_cnn.py
import torch
from torch import nn
class SimpleCNN(nn.Module):
"""Very small CNN used for a quick MNIST experiment."""
def __init__(self, num_classes: int = 10):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(1, 16, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(16, 32, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(32 * 7 * 7, 64),
nn.ReLU(),
nn.Linear(64, num_classes),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.features(x)
return self.classifier(x)
def make_model(num_classes: int = 10, device: str = "cpu"):
return SimpleCNN(num_classes=num_classes).to(device)
train.py
from pathlib import Path
import torch
from torch import nn
from torch.utils.data import DataLoader
from mnist_dataset import MNISTConfig, MNISTDataset
from simple_cnn import make_model
CHECKPOINT_PATH = Path("simple_cnn_mnist.pt")
def accuracy(model: nn.Module, data_loader, device: str) -> float:
model.eval()
correct = 0
total = 0
with torch.no_grad():
for images, labels in data_loader:
images = images.to(device)
labels = labels.to(device)
predictions = model(images).argmax(dim=1)
correct += (predictions == labels).sum().item()
total += labels.numel()
return correct / total
def train_one_fold(
train_subset,
batch_size: int,
epochs: int,
learning_rate: float,
device: str,
) -> nn.Module:
"""Train a fresh model for one K-Fold split and return it."""
model = make_model(device=device)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
loss_fn = nn.CrossEntropyLoss()
train_loader = DataLoader(train_subset, batch_size=batch_size, shuffle=True)
for _ in range(epochs):
model.train()
for images, labels in train_loader:
images = images.to(device)
labels = labels.to(device)
optimizer.zero_grad()
loss = loss_fn(model(images), labels)
loss.backward()
optimizer.step()
return model
def train(
epochs: int = 3,
learning_rate: float = 1e-3,
checkpoint_path: Path = CHECKPOINT_PATH,
) -> None:
device = "cuda" if torch.cuda.is_available() else "cpu"
mnist = MNISTDataset(MNISTConfig())
train_loader = mnist.train_loader(shuffle=True)
test_loader = mnist.test_loader()
model = make_model(device=device)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
loss_fn = nn.CrossEntropyLoss()
for epoch in range(1, epochs + 1):
model.train()
running_loss = 0.0
for images, labels in train_loader:
images = images.to(device)
labels = labels.to(device)
optimizer.zero_grad()
logits = model(images)
loss = loss_fn(logits, labels)
loss.backward()
optimizer.step()
running_loss += loss.item() * labels.size(0)
train_loss = running_loss / len(train_loader.dataset)
test_accuracy = accuracy(model, test_loader, device)
print(
f"epoch={epoch} train_loss={train_loss:.4f} "
f"test_accuracy={test_accuracy:.4f}"
)
torch.save(model.state_dict(), checkpoint_path)
print(f"saved checkpoint: {checkpoint_path}")
if __name__ == "__main__":
train()
visualizer.py
import base64
from io import BytesIO
from pathlib import Path
import numpy as np
from torchvision import datasets
def _image_to_base64(image) -> str:
buffer = BytesIO()
image.save(buffer, format="PNG")
return base64.b64encode(buffer.getvalue()).decode("ascii")
def save_suspicious_html(
indices,
labels,
pred_probs,
output_path: str | Path = "suspects.html",
data_dir: str = "data",
) -> Path:
"""Render cleanlab's suspicious MNIST examples as a small HTML gallery."""
output_path = Path(output_path)
raw_mnist = datasets.MNIST(root=data_dir, train=True, download=True)
cards = []
for rank, index in enumerate(indices, start=1):
index = int(index)
image, _ = raw_mnist[index]
given_label = int(labels[index])
predicted_label = int(np.argmax(pred_probs[index]))
confidence_given_label = float(pred_probs[index][given_label])
confidence_predicted_label = float(pred_probs[index][predicted_label])
cards.append(
f"""
<article class="card">
<div class="rank">#{rank} index={index}</div>
<img src="data:image/png;base64,{_image_to_base64(image)}" alt="MNIST sample {index}">
<dl>
<dt>Given label</dt><dd>{given_label}</dd>
<dt>Model guess</dt><dd>{predicted_label}</dd>
<dt>P(given)</dt><dd>{confidence_given_label:.4f}</dd>
<dt>P(guess)</dt><dd>{confidence_predicted_label:.4f}</dd>
</dl>
</article>
"""
)
document = f"""
<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8">
<title>MNIST cleanlab suspicious labels</title>
<style>
body {{
margin: 0;
font-family: Arial, sans-serif;
background: #f4f6f8;
color: #1f2933;
}}
header {{
padding: 28px 32px 12px;
}}
h1 {{
margin: 0;
font-size: 28px;
}}
p {{
margin: 8px 0 0;
color: #52606d;
}}
main {{
display: grid;
grid-template-columns: repeat(auto-fill, minmax(180px, 1fr));
gap: 16px;
padding: 20px 32px 32px;
}}
.card {{
border: 1px solid #d9e2ec;
border-radius: 8px;
background: white;
padding: 14px;
}}
.rank {{
font-size: 13px;
color: #52606d;
margin-bottom: 10px;
}}
img {{
display: block;
width: 112px;
height: 112px;
image-rendering: pixelated;
margin: 0 auto 12px;
border: 1px solid #bcccdc;
background: #111827;
}}
dl {{
display: grid;
grid-template-columns: 1fr auto;
gap: 6px 10px;
margin: 0;
font-size: 14px;
}}
dt {{
color: #52606d;
}}
dd {{
margin: 0;
font-weight: 700;
}}
</style>
</head>
<body>
<header>
<h1>MNIST cleanlab suspicious labels</h1>
<p>Low P(given) examples are likely annotation mistakes or ambiguous digits.</p>
</header>
<main>
{"".join(cards)}
</main>
</body>
</html>
"""
output_path.write_text(document, encoding="utf-8")
return output_path
github.com
