ML: Классификация изображений

Решение

1. Выбор архитектуры нейросети

1.1 Сравнение архитектур

Рекомендация: EfficientNet-B0 — оптимальное соотношение скорости и точности.

1.2 Архитектура сети (пример для EfficientNet)

import torch
import torch.nn as nn
from torchvision import models
from torchvision.models import EfficientNet_B0_Weights

# Загрузка предобученной модели
efficient_net = models.efficientnet_b0(weights=EfficientNet_B0_Weights.IMAGENET1K_V1)

# Заморозка ранних слоев (transfer learning)
for param in list(efficient_net.parameters())[:-20]:
    param.requires_grad = False

class ProductClassifier(nn.Module):
    def __init__(self, num_classes=10):
        super(ProductClassifier, self).__init__()
        self.backbone = efficient_net
        
        # Заменяем последний слой (был 1000 классов ImageNet)
        num_features = efficient_net.classifier[1].in_features
        
        # Добавляем голову модели
        self.classifier = nn.Sequential(
            nn.Dropout(p=0.3),
            nn.Linear(num_features, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.3),
            nn.Linear(256, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.2),
            nn.Linear(128, num_classes)
        )
        
        # Переопределяем выходной слой
        self.backbone.classifier = nn.Identity()  # убираем исходный классификатор
    
    def forward(self, x):
        features = self.backbone(x)
        logits = self.classifier(features)
        return logits

model = ProductClassifier(num_classes=10)
print(f"Модель создана с {sum(p.numel() for p in model.parameters()):,} параметров")
print(f"Обучаемых параметров: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")

2. Transfer Learning

Transfer learning позволяет использовать знания, полученные из больших датасетов (ImageNet) для своей задачи.

import torch.optim as optim
from torch.optim.lr_scheduler import ReduceLROnPlateau
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset

# === ШАГИ TRANSFER LEARNING ===

# Шаг 1: Заморозка ранних слоев (они уже знают низкоуровневые признаки)
def freeze_backbone(model, num_layers_to_freeze=20):
    params_to_freeze = list(model.backbone.parameters())[:-num_layers_to_freeze]
    for param in params_to_freeze:
        param.requires_grad = False
    
    params_to_freeze = list(model.classifier.parameters())
    for param in params_to_freeze:
        param.requires_grad = True
    
    return model

model = freeze_backbone(model)
print(f"Параметры заморожены. Обучаемых параметров: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")

# Шаг 2: Начальное обучение с заморозкой
optimizer = optim.Adam(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=0.001,
    weight_decay=1e-4
)

criterion = nn.CrossEntropyLoss()
loss_fn = criterion

# Обучаем 5 эпох с заморозкой
print("\nОбучение ЭТАП 1: Заморозка ранних слоев (5 эпох, lr=0.001)")
# ... обучение ...

# Шаг 3: Разморозка поздних слоев и fine-tuning с низким learning rate
def unfreeze_backbone(model, num_layers_to_unfreeze=10):
    params = list(model.backbone.parameters())
    
    for param in params[-num_layers_to_unfreeze:]:
        param.requires_grad = True
    
    return model

model = unfreeze_backbone(model, num_layers_to_unfreeze=10)

optimizer = optim.Adam(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=0.0001,  # намного ниже!
    weight_decay=1e-4
)

print("\nОбучение ЭТАП 2: Разморозка + Fine-tuning (10 эпох, lr=0.0001)")
# ... обучение ...

print(f"\nИтого обучаемых параметров на этапе 2: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")

# === СТРАТЕГИИ TRANSFER LEARNING ===

print("\nСтратегии Transfer Learning:")
print("""
1. FEATURE EXTRACTION (советуется для малых датасетов < 10k изображений):
   - Заморозим все слои backbone
   - Обучим только новый классификатор на вершине
   - Быстро, требует мало данных

2. FINE-TUNING (для средних датасетов 10k-100k):
   - Обучаем новый классификатор с заморозкой
   - Потом размораживаем поздние слои backbone
   - Обучаем с очень низким learning rate
   - Баланс между скоростью и точностью

3. FULL RETRAINING (для больших датасетов > 100k):
   - Инициализируем весами ImageNet
   - Обучаем всю сеть с низким lr
   - Может привести к переобучению
""")

3. Аугментация данных

Зачем? Увеличивает размер датасета и уменьшает переобучение.

import albumentations as A
from albumentations.pytorch import ToTensorV2

# === АУГМЕНТАЦИЯ ДЛЯ ОБУЧЕНИЯ ===

train_transforms = A.Compose([
    # ГЕОМЕТРИЧЕСКИЕ ПРЕОБРАЗОВАНИЯ
    A.Resize(224, 224),  # EfficientNet входной размер
    A.RandomRotate90(p=0.3),  # повороты
    A.Affine(scale=(0.9, 1.1), p=0.3),  # масштабирование
    A.Perspective(scale=(0.05, 0.1), p=0.3),  # перспектива
    A.HorizontalFlip(p=0.5),  # горизонтальное отражение
    
    # ЦВЕТОВЫЕ ТРАНСФОРМАЦИИ
    A.OneOf([
        A.GaussNoise(p=0.3),  # гауссовский шум
        A.GlassBlur(p=0.3),  # размытие
    ], p=0.2),
    
    A.OneOf([
        A.Brightness(limit=0.2, p=0.5),  # яркость
        A.Contrast(limit=0.2, p=0.5),  # контрастность
        A.RandomBrightnessContrast(p=0.5),
    ], p=0.5),
    
    A.HueSaturation(hue_shift_limit=20, sat_shift_limit=30, p=0.3),  # оттенок и насыщение
    A.RandomRain(p=0.1),  # эффекты погоды
    
    # ДОПОЛНИТЕЛЬНОЕ УЛУЧШЕНИЕ
    A.GaussianBlur(blur_limit=3, p=0.3),  # небольшое размытие
    A.Sharpen(p=0.3),  # увеличение резкости
    
    # НОРМАЛИЗАЦИЯ
    A.Normalize(
        mean=[0.485, 0.456, 0.406],  # ImageNet mean
        std=[0.229, 0.224, 0.225],   # ImageNet std
    ),
    ToTensorV2()
], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['class_labels']))  # если есть bbox

# === АУГМЕНТАЦИЯ ДЛЯ ВАЛИДАЦИИ (минимальная) ===

val_transforms = A.Compose([
    A.Resize(224, 224),
    A.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225],
    ),
    ToTensorV2()
])

print(f"Аугментация настроена")
print(f"Размер входа: 224x224 (стандарт для EfficientNet)")

# === CUSTOM DATASET ===

class ProductDataset(Dataset):
    def __init__(self, image_paths, labels, transforms=None):
        self.image_paths = image_paths
        self.labels = labels
        self.transforms = transforms
    
    def __len__(self):
        return len(self.image_paths)
    
    def __getitem__(self, idx):
        from PIL import Image
        
        image = Image.open(self.image_paths[idx]).convert('RGB')
        image = np.array(image)
        
        if self.transforms:
            augmented = self.transforms(image=image)
            image = augmented['image']
        
        label = torch.tensor(self.labels[idx], dtype=torch.long)
        
        return image, label

# Создание даталодеров
train_dataset = ProductDataset(
    image_paths=train_paths,
    labels=train_labels,
    transforms=train_transforms
)

val_dataset = ProductDataset(
    image_paths=val_paths,
    labels=val_labels,
    transforms=val_transforms
)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)

print(f"\nДатасеты загружены:")
print(f"Обучение: {len(train_dataset)} изображений")
print(f"Валидация: {len(val_dataset)} изображений")

4. Оценка качества модели

from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
from sklearn.metrics import confusion_matrix, classification_report
import matplotlib.pyplot as plt
import seaborn as sns

# === ОБУЧЕНИЕ С ОТСЛЕЖИВАНИЕМ МЕТРИК ===

def train_epoch(model, train_loader, criterion, optimizer, device):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
    
    epoch_loss = running_loss / len(train_loader)
    epoch_acc = 100 * correct / total
    
    return epoch_loss, epoch_acc

def validate(model, val_loader, criterion, device):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0
    all_preds = []
    all_labels = []
    
    with torch.no_grad():
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)
            
            outputs = model(images)
            loss = criterion(outputs, labels)
            
            running_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            
            all_preds.extend(predicted.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())
    
    epoch_loss = running_loss / len(val_loader)
    epoch_acc = 100 * correct / total
    
    return epoch_loss, epoch_acc, all_preds, all_labels

# === ГЛАВНЫЙ ЦИКЛ ОБУЧЕНИЯ ===

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = ProductClassifier(num_classes=10).to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-4)
scheduler = ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=3, verbose=True)

num_epochs = 20
best_val_acc = 0
patience = 5
patience_counter = 0

train_losses, val_losses = [], []
train_accs, val_accs = [], []

for epoch in range(num_epochs):
    train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, device)
    val_loss, val_acc, preds, labels = validate(model, val_loader, criterion, device)
    
    train_losses.append(train_loss)
    val_losses.append(val_loss)
    train_accs.append(train_acc)
    val_accs.append(val_acc)
    
    print(f"Epoch {epoch+1:2d}/{num_epochs} | Train: Loss={train_loss:.4f} Acc={train_acc:.2f}% | Val: Loss={val_loss:.4f} Acc={val_acc:.2f}%")
    
    # Early stopping
    if val_acc > best_val_acc:
        best_val_acc = val_acc
        patience_counter = 0
        torch.save(model.state_dict(), 'best_model.pth')
        print(f"  ✓ Новый лучший результат: {val_acc:.2f}%")
    else:
        patience_counter += 1
        if patience_counter >= patience:
            print(f"\nEarly stopping. Лучший результат: {best_val_acc:.2f}%")
            break
    
    scheduler.step(val_acc)

# === ФИНАЛЬНАЯ ОЦЕНКА ===

model.load_state_dict(torch.load('best_model.pth'))
_, _, final_preds, final_labels = validate(model, val_loader, criterion, device)

# Метрики
accuracy = accuracy_score(final_labels, final_preds)
precision = precision_score(final_labels, final_preds, average='weighted')
recall = recall_score(final_labels, final_preds, average='weighted')
f1 = f1_score(final_labels, final_preds, average='weighted')

print(f"\n=== ФИНАЛЬНЫЕ МЕТРИКИ ===")
print(f"Accuracy:  {accuracy:.4f}")
print(f"Precision: {precision:.4f}")
print(f"Recall:    {recall:.4f}")
print(f"F1-Score:  {f1:.4f}")

# Confusion Matrix
cm = confusion_matrix(final_labels, final_preds)
plt.figure(figsize=(10, 8))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
plt.title('Confusion Matrix')
plt.ylabel('True Label')
plt.xlabel('Predicted Label')
plt.savefig('confusion_matrix.png', dpi=150, bbox_inches='tight')
print(f"\nMatrица ошибок сохранена в confusion_matrix.png")

# Classification Report
print(f"\n=== CLASSIFICATION REPORT ===")
print(classification_report(final_labels, final_preds, target_names=[f'Class {i}' for i in range(10)]))

# Графики обучения
plt.figure(figsize=(12, 4))

plt.subplot(1, 2, 1)
plt.plot(train_losses, label='Train Loss')
plt.plot(val_losses, label='Val Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training Curves - Loss')
plt.legend()
plt.grid(True, alpha=0.3)

plt.subplot(1, 2, 2)
plt.plot(train_accs, label='Train Accuracy')
plt.plot(val_accs, label='Val Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy (%)')
plt.title('Training Curves - Accuracy')
plt.legend()
plt.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('training_curves.png', dpi=150, bbox_inches='tight')
print(f"Графики обучения сохранены в training_curves.png")

# === ИНТЕРПРЕТАЦИЯ РЕЗУЛЬТАТОВ ===

print(f"""
=== РЕКОМЕНДАЦИИ ПОСЛЕ ОБУЧЕНИЯ ===

✓ Если Accuracy > 90%:
  - Модель готова к production
  - Сохрани её с помощью torch.jit.script() для оптимизации
  - Развёртывай на Flask/FastAPI

⚠ Если Accuracy 80-90%:
  - Хороший результат, но есть место для улучшения
  - Попробуй:
    * Увеличить датасет
    * Добавить больше аугментаций
    * Использовать более мощную архитектуру (EfficientNet-B2)
    * Увеличить размер входа (256x256 вместо 224x224)

✗ Если Accuracy < 80%:
  - Есть серьёзные проблемы
  - Проверь:
    * Качество данных (есть ли дубликаты, шумные изображения?)
    * Балансировку классов
    * Корректность labels
    * Попробуй ещё более мощную модель (ResNet-50, ViT)
""")

Итоговый workflow

1. ПОДГОТОВКА ДАННЫХ
   ├─ Загрузи изображения (минимум 100 на класс)
   ├─ Проверь качество (удали битые, размытые)
   ├─ Разделить на train/val/test (70/15/15)
   └─ Примени аугментацию

2. ВЫБОР АРХИТЕКТУРЫ
   ├─ EfficientNet-B0 (рекомендуется)
   ├─ Transfer Learning из ImageNet
   └─ Кастомная голова классификатора

3. ОБУЧЕНИЕ
   ├─ ЭТАП 1: Заморозка backbone, обучение головы (5 эпох)
   ├─ ЭТАП 2: Fine-tuning поздних слоев (10 эпох, низкий LR)
   ├─ Early Stopping при отсутствии улучшения
   └─ Сохранение лучшей модели

4. ОЦЕНКА
   ├─ Accuracy, Precision, Recall, F1 на test set
   ├─ Confusion Matrix для анализа ошибок
   ├─ Per-class метрики
   └─ A/B тестирование в production

5. DEPLOYMENT
   ├─ Оптимизация модели (torch.jit, quantization)
   ├─ Контейнеризация (Docker)
   ├─ API (FastAPI + Uvicorn)
   └─ Мониторинг в production