examples/python/py-502-333-004/main.py - master - spirzen/it-code-examples

examples/python/py-502-333-004/main.py
58 строк2 KB
Spirzen
Code migration pack
09 июн 2026, 01:41
09 июн 2026, 01:41cebc284
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from torch.utils.data import DataLoader, TensorDataset

# 1. Данные и масштабирование
iris = load_iris()
X_train, X_val, y_train, y_val = train_test_split(
    iris.data, iris.target, test_size=0.2, random_state=42, stratify=iris.target
)
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_val = scaler.transform(X_val)

# 2. Тензоры и загрузчики
train_ds = TensorDataset(
    torch.tensor(X_train, dtype=torch.float32),
    torch.tensor(y_train, dtype=torch.long),
)
val_ds = TensorDataset(
    torch.tensor(X_val, dtype=torch.float32),
    torch.tensor(y_val, dtype=torch.long),
)
train_loader = DataLoader(train_ds, batch_size=16, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=16)

# 3. Модель, loss, optimizer
device = torch.device("cpu")
model = TinyNet(n_in=4, n_hidden=16, n_out=3).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-2)

# 4. Обучение с валидацией
for epoch in range(50):
    model.train()
    for batch_x, batch_y in train_loader:
        batch_x, batch_y = batch_x.to(device), batch_y.to(device)
        optimizer.zero_grad()
        loss = criterion(model(batch_x), batch_y)
        loss.backward()
        optimizer.step()

    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for batch_x, batch_y in val_loader:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)
            pred = model(batch_x).argmax(dim=1)
            correct += (pred == batch_y).sum().item()
            total += batch_y.size(0)
    print(f"epoch {epoch + 1}, val_acc={correct / total:.3f}")

# 5. Сохранение
torch.save(model.state_dict(), "iris_tiny.pt")