sandbox/mlflow-practice
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

init

Browse Source
This commit is contained in:
cs
2026-07-20 14:10:12 +03:00
commit 2c07fd63c0
10 changed files with 764 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/compare_runs.py
+54
View File
@@ -0,0 +1,54 @@
"""
Урок 5: Сравнение запусков программно
=======================================
Показывает, как через API получить все run'ы эксперимента
и сравнить их по метрикам — без веб-интерфейса.
Запуск:
python src/compare_runs.py --experiment digits_classification
"""
import argparse
import mlflow
from mlflow.tracking import MlflowClient
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--experiment", type=str, default="digits_classification")
parser.add_argument("--metric", type=str, default="accuracy")
parser.add_argument("--top", type=int, default=10)
args = parser.parse_args()
client = MlflowClient()
exp = client.get_experiment_by_name(args.experiment)
if exp is None:
print(f"❌ Эксперимент '{args.experiment}' не найден!")
return
runs = client.search_runs(
experiment_ids=[exp.experiment_id],
order_by=[f"metrics.{args.metric} DESC"],
max_results=args.top,
)
if not runs:
print("❌ Нет запусков!")
return
print(f"📊 Топ-{len(runs)} запусков в '{args.experiment}' по {args.metric}:")
print(f"{'#':>3} | {'Run ID':>36} | {'accuracy':>9} | {'n_est':>6} | {'depth':>5}")
print("-" * 75)
for i, run in enumerate(runs):
run_id = run.info.run_id
acc = run.data.metrics.get(args.metric, 0)
n_est = run.data.params.get("n_estimators", "—")
depth = run.data.params.get("max_depth", "—")
print(f"{i+1:>3} | {run_id:>36} | {acc:>9.4f} | {n_est:>6} | {depth:>5}")
best = runs[0]
print(f"\n🏆 Лучший: accuracy={best.data.metrics.get(args.metric, 0):.4f}")
if __name__ == "__main__":
main()
src/load_and_predict.py
+54
View File
@@ -0,0 +1,54 @@
"""
Урок 4: Загрузка модели из Model Registry и предсказание
==========================================================
Демонстрирует, как загрузить модель из реестра (по имени + стадия)
и сделать предсказание.
Запуск:
python src/load_and_predict.py
"""
import numpy as np
from sklearn.datasets import load_digits
import mlflow
def main():
# Загружаем модель из реестра по имени и стадии
model_name = "digits_rf_model"
stage = "Production"
print(f"📥 Загружаем модель '{model_name}' ({stage})...")
try:
model = mlflow.sklearn.load_model(
model_uri=f"models:/{model_name}/{stage}"
)
except Exception as e:
print(f"❌ Не удалось загрузить: {e}")
print(" Сначала запустите train_simple.py и register_model.py")
return
# Делаем предсказание на нескольких образцах
digits = load_digits()
X, y = digits.data, digits.target
n_samples = 5
indices = np.random.choice(len(X), n_samples, replace=False)
samples = X[indices]
true_labels = y[indices]
predictions = model.predict(samples)
print(f"\n🔮 Предсказания для {n_samples} случайных образов:")
print(f"{'Образ':>6} | {'Истинный':>8} | {'Предсказанный':>14} | {'OK':>4}")
print("-" * 45)
for i in range(n_samples):
ok = "✅" if predictions[i] == true_labels[i] else "❌"
print(f"{i+1:>6} | {true_labels[i]:>8} | {predictions[i]:>14} | {ok}")
acc = np.mean(predictions == true_labels)
print(f"\nТочность на {n_samples} образцах: {acc:.2%}")
if __name__ == "__main__":
main()
src/register_model.py
+83
View File
@@ -0,0 +1,83 @@
"""
Урок 3: Model Registry — управление версиями моделей
=====================================================
Ищет лучший run по метрике accuracy, регистрирует модель
в Model Registry и переводит её в стадию Production.
Запуск:
python src/register_model.py --experiment digits_classification
"""
import argparse
import mlflow
from mlflow.tracking import MlflowClient
def main():
parser = argparse.ArgumentParser(description="Register best model in MLflow")
parser.add_argument("--experiment", type=str, default="digits_classification")
parser.add_argument("--metric", type=str, default="accuracy")
parser.add_argument("--model-name", type=str, default="digits_rf_model")
args = parser.parse_args()
client = MlflowClient()
exp = client.get_experiment_by_name(args.experiment)
if exp is None:
print(f"❌ Эксперимент '{args.experiment}' не найден!")
print(" Сначала запустите train_simple.py")
return
# --- Ищем лучший run по метрике ---
runs = client.search_runs(
experiment_ids=[exp.experiment_id],
order_by=[f"metrics.{args.metric} DESC"],
max_results=1,
)
if not runs:
print("❌ Нет запусков в эксперименте!")
return
best_run = runs[0]
best_acc = best_run.data.metrics.get(args.metric, 0)
run_id = best_run.info.run_id
print(f"🏆 Лучший run: {run_id}")
print(f" {args.metric} = {best_acc:.4f}")
print(f" Параметры: {dict(best_run.data.params)}")
# --- Регистрируем модель ---
model_uri = f"runs:/{run_id}/model"
result = mlflow.register_model(model_uri=model_uri, name=args.model_name)
print(f"\n📦 Модель зарегистрирована: {args.model_name} v{result.version}")
# --- Переводим в Production ---
client = MlflowClient()
client.set_registered_model_tag(args.model_name, "task", "classification")
# Снимаем Production со всех предыдущих версий
versions = client.search_model_versions(f"name='{args.model_name}'")
for v in versions:
if v.current_stage == "Production" and v.version != result.version:
client.transition_model_version_stage(
name=args.model_name,
version=v.version,
stage="Archived",
)
print(f" ↓ Версия {v.version} → Archived")
client.transition_model_version_stage(
name=args.model_name,
version=result.version,
stage="Production",
archive_existing_versions=True,
)
print(f" ↑ Версия {result.version} → Production ✅")
# --- Список всех версий ---
print(f"\n📋 Все версии модели '{args.model_name}':")
versions = client.search_model_versions(f"name='{args.model_name}'")
for v in versions:
print(f" v{v.version} | stage={v.current_stage} | run_id={v.run_id}")
if __name__ == "__main__":
main()
src/train_gpu.py
+207
View File
@@ -0,0 +1,207 @@
"""
Урок 2: MLflow + PyTorch на GPU (RTX 3090)
=============================================
Обучаем CNN на MNIST. Автоматически использует CUDA если доступна,
иначе — CPU. Все параметры, метрики и модель логируются в MLflow.
Запуск на 3090:
python src/train_gpu.py --epochs 10 --batch-size 256 --lr 0.001
Запуск на CPU (для теста):
python src/train_gpu.py --epochs 2 --batch-size 64 --lr 0.01
"""
import argparse
import os
import time
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import torch
from sklearn.metrics import recall_score, precision_score, f1_score
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import mlflow
import mlflow.pytorch
# ─── Модель: простая CNN ───
class SimpleCNN(nn.Module):
def __init__(self, num_classes=10):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
)
self.classifier = nn.Sequential(
nn.Linear(64 * 7 * 7, 128),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(128, num_classes),
)
def forward(self, x):
x = self.features(x)
x = x.view(x.size(0), -1)
return self.classifier(x)
def main():
parser = argparse.ArgumentParser(description="MLflow + PyTorch GPU demo")
parser.add_argument("--epochs", type=int, default=5)
parser.add_argument("--batch-size", type=int, default=128)
parser.add_argument("--lr", type=float, default=0.001)
parser.add_argument("--experiment-name", type=str, default="mnist_cnn_gpu")
args = parser.parse_args()
# ─── Устройство ───
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"🖥️ Устройство: {device}")
if device.type == "cuda":
gpu_name = torch.cuda.get_device_name(0)
gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
print(f" GPU: {gpu_name} ({gpu_mem:.1f} GB)")
# ─── MLflow эксперимент ───
mlflow.set_experiment(args.experiment_name)
# ─── Данные ───
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,)),
])
data_dir = "data"
os.makedirs(data_dir, exist_ok=True)
train_ds = datasets.MNIST(data_dir, train=True, download=True, transform=transform)
test_ds = datasets.MNIST(data_dir, train=False, download=True, transform=transform)
train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, num_workers=4)
test_loader = DataLoader(test_ds, batch_size=args.batch_size, shuffle=False, num_workers=4)
print(f"📊 Train: {len(train_ds)}, Test: {len(test_ds)}")
# ─── MLflow run ───
with mlflow.start_run(run_name=f"cnn_e{args.epochs}_bs{args.batch_size}") as run:
print(f"MLflow run ID: {run.info.run_id}")
# Логируем параметры
mlflow.log_param("epochs", args.epochs)
mlflow.log_param("batch_size", args.batch_size)
mlflow.log_param("learning_rate", args.lr)
mlflow.log_param("device", str(device))
mlflow.log_param("optimizer", "Adam")
if device.type == "cuda":
mlflow.log_param("gpu_name", gpu_name)
# ─── Модель, оптимизатор, функция потерь ───
model = SimpleCNN().to(device)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
criterion = nn.CrossEntropyLoss()
# ─── Цикл обучения ───
train_losses, test_accs = [], []
for epoch in range(args.epochs):
model.train()
epoch_loss = 0.0
t0 = time.time()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
if batch_idx % 100 == 0:
print(f" Epoch {epoch+1}/{args.epochs} "
f"[{batch_idx * len(data)}/{len(train_loader.dataset)}] "
f"loss={loss.item():.4f}")
avg_loss = epoch_loss / len(train_loader)
train_losses.append(avg_loss)
# ─── Валидация ───
model.eval()
correct, total = 0, 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
pred = output.argmax(dim=1)
correct += (pred == target).sum().item()
total += target.size(0)
acc = correct / total
test_accs.append(acc)
# Recall / Precision / F1 (macro-усреднение по 10 классам)
all_preds, all_labels = [], []
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
pred = output.argmax(dim=1)
all_preds.extend(pred.cpu().numpy())
all_labels.extend(target.cpu().numpy())
recall = recall_score(all_labels, all_preds, average="macro")
precision = precision_score(all_labels, all_preds, average="macro")
f1 = f1_score(all_labels, all_preds, average="macro")
elapsed = time.time() - t0
# Логируем метрики по эпохам
mlflow.log_metric("train_loss", avg_loss, step=epoch)
mlflow.log_metric("test_accuracy", acc, step=epoch)
mlflow.log_metric("recall_macro", recall, step=epoch)
mlflow.log_metric("precision_macro", precision, step=epoch)
mlflow.log_metric("f1_macro", f1, step=epoch)
mlflow.log_metric("epoch_time_sec", elapsed, step=epoch)
print(f" → Epoch {epoch+1}: loss={avg_loss:.4f}, "
f"acc={acc:.4f}, time={elapsed:.1f}s")
# ─── График обучения ───
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
ax1.plot(range(1, args.epochs + 1), train_losses, "b-o")
ax1.set_title("Train Loss")
ax1.set_xlabel("Epoch")
ax2.plot(range(1, args.epochs + 1), test_accs, "g-o")
ax2.set_title("Test Accuracy")
ax2.set_xlabel("Epoch")
plt.tight_layout()
os.makedirs("artifacts", exist_ok=True)
plot_path = "artifacts/training_curve.png"
fig.savefig(plot_path)
mlflow.log_artifact(plot_path)
plt.close(fig)
# ─── Логируем модель ───
mlflow.pytorch.log_model(
model,
artifact_path="model",
serialization_format="pickle",
registered_model_name=None,
input_example=torch.randn(1, 1, 28, 28).to(device),
)
# ─── Теги ───
mlflow.set_tag("model_type", "SimpleCNN")
mlflow.set_tag("framework", "PyTorch")
mlflow.set_tag("dataset", "MNIST")
final_acc = test_accs[-1]
print(f"\n✅ Обучение завершено!")
print(f" Финальная accuracy: {final_acc:.4f}")
print(f" Run ID: {run.info.run_id}")
print(f" Артефакты: training_curve.png, model/")
if __name__ == "__main__":
main()
src/train_simple.py
+115
View File
@@ -0,0 +1,115 @@
"""
Урок 1: Основы MLflow Tracking
================================
Обучаем RandomForest на встроенном датасете digits (scikit-learn).
Логируем параметры, метрики, модель и графики в MLflow.
Запуск:
python src/train_simple.py --n-estimators 100 --max-depth 8
"""
import argparse
import os
import matplotlib
matplotlib.use("Agg") # без GUI
import matplotlib.pyplot as plt
import numpy as np
from sklearn.datasets import load_digits
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, confusion_matrix, classification_report, recall_score, precision_score, f1_score
from sklearn.model_selection import train_test_split
import mlflow
import mlflow.sklearn
def main():
parser = argparse.ArgumentParser(description="MLflow + scikit-learn demo")
parser.add_argument("--n-estimators", type=int, default=100, help="количество деревьев")
parser.add_argument("--max-depth", type=int, default=8, help="макс. глубина дерева")
parser.add_argument("--experiment-name", type=str, default="digits_classification")
args = parser.parse_args()
# --- MLflow: задаём эксперимент ---
mlflow.set_experiment(args.experiment_name)
# --- Данные ---
digits = load_digits()
X, y = digits.data, digits.target
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42
)
print(f"Датасет digits: {X.shape[0]} образцов, {X.shape[1]} признаков")
# --- MLflow: начинаем run ---
with mlflow.start_run(run_name=f"rf_{args.n_estimators}_{args.max_depth}") as run:
print(f"MLflow run ID: {run.info.run_id}")
# Логируем параметры (до обучения!)
mlflow.log_param("n_estimators", args.n_estimators)
mlflow.log_param("max_depth", args.max_depth)
mlflow.log_param("dataset", "digits")
mlflow.log_param("test_size", 0.2)
# --- Обучение ---
model = RandomForestClassifier(
n_estimators=args.n_estimators,
max_depth=args.max_depth,
random_state=42,
n_jobs=-1,
)
model.fit(X_train, y_train)
# --- Предсказание и метрики ---
y_pred = model.predict(X_test)
acc = accuracy_score(y_test, y_pred)
print(f"Accuracy: {acc:.4f}")
# Логируем метрики
mlflow.log_metric("accuracy", acc)
mlflow.log_metric("recall_macro", recall_score(y_test, y_pred, average="macro"))
mlflow.log_metric("precision_macro", precision_score(y_test, y_pred, average="macro"))
mlflow.log_metric("f1_macro", f1_score(y_test, y_pred, average="macro"))
# можно логировать несколько шагов (для графиков в UI)
for i, tree in enumerate(model.estimators_):
tree_acc = accuracy_score(y_test, tree.predict(X_test))
mlflow.log_metric("per_tree_accuracy", tree_acc, step=i)
# --- Confusion matrix как артефакт ---
cm = confusion_matrix(y_test, y_pred)
fig, ax = plt.subplots(figsize=(8, 6))
ax.imshow(cm, cmap="Blues")
ax.set_title("Confusion Matrix")
ax.set_xlabel("Predicted")
ax.set_ylabel("Actual")
plt.tight_layout()
os.makedirs("artifacts", exist_ok=True)
cm_path = "artifacts/confusion_matrix.png"
fig.savefig(cm_path)
mlflow.log_artifact(cm_path)
plt.close(fig)
# --- Classification report как текстовый артефакт ---
report = classification_report(y_test, y_pred)
report_path = "artifacts/classification_report.txt"
with open(report_path, "w") as f:
f.write(report)
mlflow.log_artifact(report_path)
# --- Логируем саму модель ---
mlflow.sklearn.log_model(
model,
artifact_path="model",
registered_model_name=None, # регистрация — в отдельном скрипте
)
# --- Теги ---
mlflow.set_tag("model_type", "RandomForest")
mlflow.set_tag("author", "practice")
print(f"\n✅ Готово! Откройте MLflow UI и найдите эксперимент '{args.experiment_name}'")
print(f" Run ID: {run.info.run_id}")
print(f" Accuracy: {acc:.4f}")
if __name__ == "__main__":
main()