Transfer Learning та візуалізація внутрішніх представлень CNN

import torch
import torch.nn as nn
import torchvision.models as models

# Завантаження pre-trained ResNet18
resnet18 = models.resnet18(weights='IMAGENET1K_V1')

print("Архітектура ResNet18:")
print(resnet18)
print(f"\nЗагальна кількість параметрів: {sum(p.numel() for p in resnet18.parameters()):,}")

# Аналіз структури
print(f"\nОстанній шар (FC): {resnet18.fc}")
print(f"Параметрів у FC шарі: {sum(p.numel() for p in resnet18.fc.parameters()):,}")
print(f"Параметрів у згорткових шарах: {sum(p.numel() for p in resnet18.parameters()) - sum(p.numel() for p in resnet18.fc.parameters()):,}")

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision
from torchvision import datasets, transforms, models
import matplotlib.pyplot as plt
import numpy as np

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.manual_seed(42)

print(f"Використовується пристрій: {device}")

# Статистика ImageNet (ResNet навчалася на цих значеннях!)
imagenet_mean = (0.485, 0.456, 0.406)
imagenet_std = (0.229, 0.224, 0.225)

# Трансформації: resize до 224×224 + нормалізація ImageNet
transform_train = transforms.Compose([
    transforms.Resize(224),  # CIFAR-10: 32×32 → 224×224
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomCrop(224, padding=16),  # Аугментація після resize
    transforms.ToTensor(),
    transforms.Normalize(imagenet_mean, imagenet_std)
])

transform_test = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(imagenet_mean, imagenet_std)
])

# Завантаження CIFAR-10
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4, pin_memory=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False, num_workers=4, pin_memory=True)

print(f"Тренувальних батчів: {len(train_loader)}")
print(f"Тестових батчів: {len(test_loader)}")

# Перевірка розміру
sample_img, _ = train_dataset[0]
print(f"Розмір зображення після трансформацій: {sample_img.shape}")

# Завантаження pre-trained ResNet18
model = models.resnet18(weights='IMAGENET1K_V1')

# Крок 1: Заморожування всіх шарів
for param in model.parameters():
    param.requires_grad = False

# Крок 2: Заміна останнього FC шару
# Оригінальний: Linear(512, 1000)
# Новий: Linear(512, 10)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 10)  # Лише цей шар буде навчатися!

model = model.to(device)

# Підрахунок навчаємих параметрів
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
frozen_params = sum(p.numel() for p in model.parameters() if not p.requires_grad)

print("="*70)
print("Transfer Learning: Feature Extraction Mode")
print("="*70)
print(f"Загальна кількість параметрів: {trainable_params + frozen_params:,}")
print(f"Навчаємих параметрів (FC шар): {trainable_params:,}")
print(f"Заморожених параметрів (згорткові шари): {frozen_params:,}")
print(f"Відсоток навчаємих параметрів: {100 * trainable_params / (trainable_params + frozen_params):.2f}%")
print("="*70)

# Оптимізатор — передаємо лише параметри FC шару
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.fc.parameters(), lr=0.001)  # Лише model.fc!
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)

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

print("\nПочаток навчання (Feature Extraction)...\n")

for epoch in range(num_epochs):
    # Тренування
    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, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
    
    train_loss = running_loss / len(train_loader)
    train_acc = 100 * correct / total
    train_losses.append(train_loss)
    train_accs.append(train_acc)
    
    # Валідація
    model.eval()
    val_loss = 0.0
    correct = 0
    total = 0
    
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            
            val_loss += loss.item()
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    
    val_loss /= len(test_loader)
    val_acc = 100 * correct / total
    val_losses.append(val_loss)
    val_accs.append(val_acc)
    
    scheduler.step()
    
    print(f"Епоха [{epoch+1}/{num_epochs}] "
          f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}% | "
          f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2f}%")

print(f"\n✅ Навчання завершено! Фінальна Val Accuracy: {val_accs[-1]:.2f}%")
print(f"🚀 Порівняйте: з нуля за 30 епох → 81%, Transfer Learning за 10 епох → {val_accs[-1]:.2f}%")

# Функція для витягування Feature Maps з проміжних шарів
def get_feature_maps(model, image, layer_name):
    """
    Отримує Feature Maps із конкретного шару моделі.
    
    Args:
        model: CNN модель
        image: Вхідне зображення (1, C, H, W)
        layer_name: Назва шару (наприклад, 'layer1', 'layer2')
    
    Returns:
        Feature Maps як тензор
    """
    feature_maps = {}
    
    def hook_fn(module, input, output):
        feature_maps['output'] = output
    
    # Реєструємо hook на потрібному шарі
    layer = dict(model.named_modules())[layer_name]
    hook = layer.register_forward_hook(hook_fn)
    
    # Прямий прохід
    model.eval()
    with torch.no_grad():
        _ = model(image)
    
    hook.remove()
    return feature_maps['output']

# Візуалізація Feature Maps для одного зображення
def visualize_feature_maps(image, layer_name, num_maps=16):
    """Візуалізує перші num_maps Feature Maps із шару"""
    
    # Отримуємо Feature Maps
    feature_maps = get_feature_maps(model, image.unsqueeze(0).to(device), layer_name)
    feature_maps = feature_maps.squeeze(0).cpu()  # (C, H, W)
    
    num_channels = min(num_maps, feature_maps.shape[0])
    fig, axes = plt.subplots(4, 4, figsize=(14, 14))
    fig.suptitle(f'Feature Maps із {layer_name} ({feature_maps.shape[0]} каналів)',
                 fontsize=16, fontweight='bold')
    
    for i, ax in enumerate(axes.flat):
        if i < num_channels:
            fmap = feature_maps[i].numpy()
            ax.imshow(fmap, cmap='viridis')
            ax.set_title(f'Канал {i+1}', fontsize=10)
            ax.axis('off')
        else:
            ax.axis('off')
    
    plt.tight_layout()
    plt.show()

# Вибираємо зображення коня з тестового набору
test_image, test_label = test_dataset[42]  # Індекс 42 — кінь
classes = ['літак', 'автомобіль', 'птах', 'кіт', 'олень',
           'собака', 'жаба', 'кінь', 'корабель', 'вантажівка']

# Візуалізація оригіналу (денормалізація)
img_display = test_image.clone()
for t, m, s in zip(img_display, imagenet_mean, imagenet_std):
    t.mul_(s).add_(m)
img_display = torch.clamp(img_display, 0, 1).permute(1, 2, 0).numpy()

plt.figure(figsize=(6, 6))
plt.imshow(img_display)
plt.title(f'Оригінальне зображення: {classes[test_label]}', fontsize=14, fontweight='bold')
plt.axis('off')
plt.show()

print(f"Клас зображення: {classes[test_label]}")
print(f"Форма вхідного зображення: {test_image.shape}")

# Візуалізація Feature Maps з різних шарів ResNet18
print("Feature Maps із layer1 (ранній шар, низькорівневі ознаки):")
visualize_feature_maps(test_image, 'layer1', num_maps=16)

print("\nFeature Maps із layer2 (середні ознаки):")
visualize_feature_maps(test_image, 'layer2', num_maps=16)

print("\nFeature Maps із layer3 (високорівневі ознаки):")
visualize_feature_maps(test_image, 'layer3', num_maps=16)

# Отримання ваг першого згорткового шару
first_conv_layer = model.conv1
filters = first_conv_layer.weight.data.cpu()  # (64, 3, 7, 7)

print(f"Форма фільтрів: {filters.shape}")
print(f"  - Кількість фільтрів: {filters.shape[0]}")
print(f"  - Вхідні канали (RGB): {filters.shape[1]}")
print(f"  - Розмір kernel: {filters.shape[2]}×{filters.shape[3]}")

# Нормалізація для візуалізації
filters_min, filters_max = filters.min(), filters.max()
filters_normalized = (filters - filters_min) / (filters_max - filters_min)

# Візуалізація перших 16 фільтрів
fig, axes = plt.subplots(4, 4, figsize=(12, 12))
fig.suptitle('Фільтри першого згорткового шару ResNet18\n(Навчені на ImageNet)', 
             fontsize=16, fontweight='bold')

for i, ax in enumerate(axes.flat):
    # Перетворюємо (3, 7, 7) → (7, 7, 3) для matplotlib
    filter_img = filters_normalized[i].permute(1, 2, 0).numpy()
    ax.imshow(filter_img)
    ax.set_title(f'Фільтр {i+1}', fontsize=10)
    ax.axis('off')

plt.tight_layout()
plt.show()

class GradCAM:
    """
    Реалізація Grad-CAM для візуалізації областей уваги CNN.
    """
    
    def __init__(self, model, target_layer):
        self.model = model
        self.target_layer = target_layer
        self.gradients = None
        self.activations = None
        
        # Реєструємо hooks
        target_layer.register_forward_hook(self.save_activation)
        target_layer.register_backward_hook(self.save_gradient)
    
    def save_activation(self, module, input, output):
        """Hook для збереження активацій"""
        self.activations = output.detach()
    
    def save_gradient(self, module, grad_input, grad_output):
        """Hook для збереження градієнтів"""
        self.gradients = grad_output[0].detach()
    
    def generate_cam(self, image, class_idx=None):
        """
        Генерує Grad-CAM для конкретного класу.
        
        Args:
            image: Вхідне зображення (1, C, H, W)
            class_idx: Індекс класу (якщо None — передбачений клас)
        
        Returns:
            cam: Теплова карта (H, W)
        """
        self.model.eval()
        
        # Прямий прохід
        output = self.model(image)
        
        if class_idx is None:
            class_idx = output.argmax(dim=1).item()
        
        # Обнуляємо градієнти та робимо backward для обраного класу
        self.model.zero_grad()
        class_loss = output[0, class_idx]
        class_loss.backward()
        
        # Обчислюємо ваги каналів (Global Average Pooling градієнтів)
        weights = self.gradients.mean(dim=(2, 3), keepdim=True)  # (1, C, 1, 1)
        
        # Зважена сума Feature Maps
        cam = (weights * self.activations).sum(dim=1, keepdim=True)  # (1, 1, H, W)
        
        # ReLU (відкидаємо від'ємні значення) та нормалізація
        cam = torch.relu(cam)
        cam = cam.squeeze().cpu().numpy()
        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
        
        return cam, class_idx

# Функція для накладання теплової карти на зображення
def overlay_heatmap(image, cam, alpha=0.5):
    """Накладає Grad-CAM на оригінальне зображення"""
    import cv2
    
    # Resize CAM до розміру зображення
    cam_resized = cv2.resize(cam, (image.shape[1], image.shape[0]))
    
    # Конвертуємо CAM у колірну теплову карту (червоний=важливо, синій=неважливо)
    heatmap = cv2.applyColorMap(np.uint8(255 * cam_resized), cv2.COLORMAP_JET)
    heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
    
    # Накладаємо з прозорістю
    overlay = (heatmap * alpha + image * 255 * (1 - alpha)).astype(np.uint8)
    
    return overlay

print("✅ Grad-CAM клас визначено")

# Візуалізація Grad-CAM для кількох зображень
grad_cam = GradCAM(model, target_layer=model.layer4[-1])  # Останній residual блок

# Вибираємо 6 зображень різних класів
test_indices = [42, 123, 234, 456, 678, 890]
classes = ['літак', 'автомобіль', 'птах', 'кіт', 'олень',
           'собака', 'жаба', 'кінь', 'корабель', 'вантажівка']

fig, axes = plt.subplots(3, 4, figsize=(18, 14))
fig.suptitle('Grad-CAM: Області уваги моделі при класифікації', fontsize=18, fontweight='bold')

for row, idx in enumerate(test_indices[:3]):
    # Завантаження зображення
    image_tensor, true_label = test_dataset[idx]
    
    # Денормалізація для візуалізації
    image_display = image_tensor.clone()
    for t, m, s in zip(image_display, imagenet_mean, imagenet_std):
        t.mul_(s).add_(m)
    image_display = torch.clamp(image_display, 0, 1).permute(1, 2, 0).numpy()
    
    # Генерація Grad-CAM
    cam, pred_class = grad_cam.generate_cam(image_tensor.unsqueeze(0).to(device))
    
    # Накладання теплової карти
    overlay = overlay_heatmap(image_display, cam, alpha=0.4)
    
    # Візуалізація
    # Колонка 1: Оригінал
    axes[row, 0].imshow(image_display)
    axes[row, 0].set_title(f'Оригінал\nСправжній клас: {classes[true_label]}', 
                           fontsize=11, fontweight='bold')
    axes[row, 0].axis('off')
    
    # Колонка 2: Теплова карта
    axes[row, 1].imshow(cam, cmap='jet')
    axes[row, 1].set_title(f'Grad-CAM\nПередбачено: {classes[pred_class]}', 
                           fontsize=11, fontweight='bold')
    axes[row, 1].axis('off')
    
    # Колонка 3: Накладання
    axes[row, 2].imshow(overlay)
    axes[row, 2].set_title('Накладання\n(червоне = важливо)', 
                           fontsize=11, fontweight='bold')
    axes[row, 2].axis('off')
    
    # Колонка 4: Топ-3 передбачення
    with torch.no_grad():
        output = model(image_tensor.unsqueeze(0).to(device))
        probs = torch.softmax(output, dim=1)[0]
        top3_probs, top3_indices = torch.topk(probs, 3)
    
    top3_text = "Топ-3 передбачення:\n"
    for i, (prob, idx_class) in enumerate(zip(top3_probs, top3_indices)):
        marker = "✅" if idx_class == true_label else "  "
        top3_text += f"{marker} {classes[idx_class]}: {prob*100:.1f}%\n"
    
    axes[row, 3].text(0.1, 0.5, top3_text, fontsize=12, verticalalignment='center',
                      family='monospace', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.3))
    axes[row, 3].axis('off')

plt.tight_layout()
plt.show()