train_data.py

import time
import torch

def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs=20, device='cpu'):
    best_val_loss = float('inf')
    model.to(device)

    for epoch in range(num_epochs):
        model.train()
        train_loss = 0.0
        start_time = time.time()

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

            optimizer.zero_grad()
            output = model(X_batch)
            loss = criterion(output, y_batch)
            loss.backward()
            optimizer.step()

            train_loss += loss.item() * X_batch.size(0)

        train_loss /= len(train_loader.dataset)

        # Validation phase
        model.eval()
        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)
                output = model(X_val)
                loss = criterion(output, y_val)
                val_loss += loss.item() * X_val.size(0)

        val_loss /= len(val_loader.dataset)
        end_time = time.time()

        print(f"Epoch {epoch+1}/{num_epochs} | "
              f"Train Loss: {train_loss:.4f} | "
              f"Val Loss: {val_loss:.4f} | "
              f"Time: {end_time - start_time:.2f} sec")

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save(model.state_dict(), 'best_transformer_model.pth')
            print("Saved new best model")

    print("Training complete.")