task2/model.py - master - akp1n/ADVML

task2/model.py
190 строк6 KB
Artemiy
fix
24 фев 2025, 20:45
24 фев 2025, 20:4557cd6cd
from typing import Tuple, List, Optional
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np
from torch.utils.data import TensorDataset, random_split, DataLoader
from sklearn.datasets import make_regression
from sklearn.decomposition import PCA
from sklearn.preprocessing import MinMaxScaler
from torch import Tensor

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

random_seed: int = 42
generator = torch.Generator().manual_seed(random_seed)

class Net(nn.Module):
    def __init__(self, input_size: int, output_size: int, is_batch_norm: bool = False, is_dropout: bool = False):
        super(Net, self).__init__()
        self.l1: nn.Linear = nn.Linear(input_size, 512)
        self.bn1: nn.BatchNorm1d = nn.BatchNorm1d(512)
        self.l2: nn.Linear = nn.Linear(512, 128)
        self.bn2: nn.BatchNorm1d = nn.BatchNorm1d(128)
        self.l3: nn.Linear = nn.Linear(128, output_size)
        self.act: nn.ReLU = nn.ReLU()
        self.dropout: nn.Dropout = nn.Dropout(p=0.5)
        self.is_batch_norm = is_batch_norm
        self.is_dropout = is_dropout

    def forward(self, x: Tensor) -> Tensor:
        x = self.l1(x)
        if self.is_batch_norm:
            x = self.bn1(x)
        x = self.act(x)
        if self.is_dropout:
            x = self.dropout(x)

        x = self.l2(x)
        if self.is_batch_norm:
            x = self.bn2(x)
        x = self.act(x)
        if self.is_dropout:
            x = self.dropout(x)

        x = self.l3(x)
        return x

def generate_data_simple(n_samples: int = 1000, noise: float = 0.1) -> List[Tuple[np.ndarray, np.ndarray]]:
    x_sin = np.linspace(-3, 3, n_samples).reshape(-1, 1)
    y_sin = np.sin(x_sin) + np.random.normal(0, noise, size=x_sin.shape)

    x_cubic = np.linspace(-3, 3, n_samples).reshape(-1, 1)
    y_cubic = np.power(x_cubic, 3) + np.random.normal(0, noise * 10, size=x_cubic.shape)

    x_tanh = np.linspace(-3, 3, n_samples).reshape(-1, 1)
    y_tanh = np.tanh(x_tanh) + np.random.normal(0, noise, size=x_tanh.shape)

    scaler = MinMaxScaler(feature_range=(-1, 1))

    y_sin = scaler.fit_transform(y_sin)
    y_cubic = scaler.fit_transform(y_cubic)
    y_tanh = scaler.fit_transform(y_tanh)

    return [(x_sin, y_sin), (x_cubic, y_cubic), (x_tanh, y_tanh)]

def train_model_with_scheduler(
    model: nn.Module,
    train_loader: DataLoader,
    val_loader: DataLoader,
    epochs: int,
    loss_fn: nn.Module,
    optimizer,
    scheduler=None
):
    model.to(device)
    train_losses = []
    val_losses = []
    lrs = []

    for epoch in range(epochs):
        model.train()
        epoch_train_loss = 0.0

        for x_batch, y_batch in train_loader:
            x_batch = x_batch.to(device)
            y_batch = y_batch.to(device)

            optimizer.zero_grad()
            y_pred = model(x_batch)
            loss = loss_fn(y_pred, y_batch)
            loss.backward()
            optimizer.step()
            epoch_train_loss += loss.item()

        if scheduler is not None:
            scheduler.step()

        epoch_train_loss /= len(train_loader)
        train_losses.append(epoch_train_loss)

        current_lr = optimizer.param_groups[0]['lr']
        lrs.append(current_lr)

        model.eval()
        epoch_val_loss = 0.0
        with torch.no_grad():
            for x_val, y_val in val_loader:
                x_val = x_val.to(device)
                y_val = y_val.to(device)
                y_pred_val = model(x_val)
                loss_val = loss_fn(y_pred_val, y_val)
                epoch_val_loss += loss_val.item()
        epoch_val_loss /= len(val_loader)
        val_losses.append(epoch_val_loss)

        print(f"Epoch [{epoch+1}/{epochs}] | Train Loss: {epoch_train_loss:.4f} | Val Loss: {epoch_val_loss:.4f} | LR: {current_lr:.6f}")

    return train_losses, val_losses, lrs

def visualize_loss_curves(train_losses, val_losses, lrs, optimizer_name, scheduler_name):
    fig, ax1 = plt.subplots(figsize=(10, 6))

    ax1.set_xlabel("Epoch")
    ax1.set_ylabel("MSE Loss", color="tab:blue")
    ax1.plot(train_losses, label="Train Loss", color="tab:blue")
    ax1.plot(val_losses, label="Val Loss", color="tab:orange")
    ax1.tick_params(axis='y', labelcolor="tab:blue")
    ax1.legend(loc='upper left')

    ax2 = ax1.twinx()
    ax2.set_ylabel("Learning Rate", color="tab:green")
    ax2.plot(lrs, label="Learning Rate", color="tab:green", linestyle='--')
    ax2.tick_params(axis='y', labelcolor="tab:green")
    ax2.legend(loc='upper right')

    plt.title(f"Loss and Learning Rate Curve: {optimizer_name} + {scheduler_name}")
    fig.tight_layout()
    filename = f"loss_lr_curve_{optimizer_name}_{scheduler_name}.png"
    plt.savefig(filename, format="png")
    print(f"График сохранён в файл: {filename}")
    plt.close()

def run_experiments_with_optimizers_and_schedulers(
    data: List[Tuple[np.ndarray, np.ndarray]],
    optimizers_list: List,
    schedulers_list: List,
    epochs=10,
    batch_size=64
):
    x_np, y_np = data[0]
    x = torch.tensor(x_np, dtype=torch.float32)
    y = torch.tensor(y_np, dtype=torch.float32)

    dataset = TensorDataset(x, y)
    train_size = int(0.8 * len(dataset))
    val_size = len(dataset) - train_size

    train_dataset, val_dataset = random_split(dataset, [train_size, val_size], generator=generator)

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

    input_size = x.shape[1]
    output_size = y.shape[1] if len(y.shape) > 1 else 1
    loss_fn = nn.MSELoss()

    for opt_class in optimizers_list:
        for sch_class in schedulers_list:
            model = Net(input_size=input_size, output_size=output_size)
            optimizer = opt_class(model.parameters(), lr=0.01)
            scheduler = sch_class(optimizer)

            print(f"\n=== Обучение: {opt_class.__name__} + {sch_class.__name__} ===")
            train_losses, val_losses, lrs = train_model_with_scheduler(
                model=model,
                train_loader=train_loader,
                val_loader=val_loader,
                epochs=epochs,
                loss_fn=loss_fn,
                optimizer=optimizer,
                scheduler=scheduler
            )

            visualize_loss_curves(
                train_losses,
                val_losses,
                lrs,
                optimizer_name=opt_class.__name__,
                scheduler_name=sch_class.__name__,
            )