model_train.py

import argparse
import os
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import (
    f1_score, accuracy_score, roc_auc_score,
    average_precision_score, multilabel_confusion_matrix,
    roc_curve, auc
)
from utils import (
    evaluate_model,
    plot_training_metrics,
    plot_combined_confusion_matrices,
    plot_roc_curves
)
from torchvision.models import resnet18, ResNet18_Weights
from torchvision.models import efficientnet_b0, EfficientNet_B0_Weights

from data_utils import get_data_loaders  # assumes your dataloader is modularized in data_utils.py

def get_model(model_name: str, num_classes: int):
    """
    Load a model architecture with modified final layers for multi-label classification.
    """
    if model_name == 'resnet18':
        model = resnet18(weights=ResNet18_Weights.DEFAULT)
        model.fc = nn.Sequential(
            nn.Linear(model.fc.in_features, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, num_classes)
        )
    elif model_name == 'efficientnet_b0':
        model = efficientnet_b0(weights=EfficientNet_B0_Weights.DEFAULT)
        model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
    else:
        raise ValueError(f"Unsupported model: {model_name}")
    return model

def get_optimizer(name, model, lr):
    """
    Return the specified optimizer.
    """
    if name == 'sgd':
        return optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
    elif name == 'adam':
        return optim.Adam(model.parameters(), lr=lr)
    elif name == 'radam':
        from torch_optimizer import RAdam
        return RAdam(model.parameters(), lr=lr)
    else:
        raise ValueError(f"Unsupported optimizer: {name}")

def train_model(args):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    os.makedirs(args.save_dir, exist_ok=True)
    train_loader, val_loader = get_data_loaders(args.data_dir, args.batch_size, args.img_size)

    model = get_model(args.model_name, num_classes=5).to(device)
    optimizer = get_optimizer(args.optimizer, model, args.lr)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1) if args.scheduler else None
    criterion = nn.BCEWithLogitsLoss()

    train_stats = {k: [] for k in ['f1', 'accuracy', 'mAP', 'AUC']}
    val_stats = {k: [] for k in ['f1', 'accuracy', 'mAP', 'AUC']}
    class_names = ["defect_0", "defect_1", "defect_2", "defect_3", "defect_4"]

    for epoch in range(args.epochs):
        model.train()
        all_labels, all_preds, all_probs = [], [], []

        for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{args.epochs}"):
            images, labels = images.to(device), labels.to(device).float()
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            probs = torch.sigmoid(outputs).detach().cpu().numpy()
            preds = (probs >= args.threshold).astype(int)

            all_labels.append(labels.cpu().numpy())
            all_preds.append(preds)
            all_probs.append(probs)

        if scheduler:
            scheduler.step()

        all_labels = np.vstack(all_labels)
        all_preds = np.vstack(all_preds)
        all_probs = np.vstack(all_probs)

        train_stats['f1'].append(f1_score(all_labels, all_preds, average='macro'))
        train_stats['accuracy'].append(accuracy_score(all_labels, all_preds))
        train_stats['AUC'].append(roc_auc_score(all_labels, all_probs, average='macro'))
        train_stats['mAP'].append(average_precision_score(all_labels, all_probs, average='macro'))

        val_metrics = evaluate_model(model, val_loader, device, criterion, args.threshold)
        for key in val_stats:
            val_stats[key].append(val_metrics[key])

        print(f"\nEpoch {epoch+1}: Train F1={train_stats['f1'][-1]:.4f}, Val F1={val_stats['f1'][-1]:.4f}")

    # Save model and plots
    torch.save(model.state_dict(), os.path.join(args.save_dir, "model.pth"))
    plot_training_metrics(train_stats, val_stats, os.path.join(args.save_dir, "training_metrics.png"))
    plot_combined_confusion_matrices(all_preds, all_labels, val_metrics["preds"], val_metrics["labels"], class_names, os.path.join(args.save_dir, "conf_matrix.png"))
    plot_roc_curves(val_metrics["labels"], val_metrics["probs"], class_names, os.path.join(args.save_dir, "roc_curve.png"))

    print(f"\nTraining complete. Artifacts saved to: {args.save_dir}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_dir', type=str, required=True)
    parser.add_argument('--save_dir', type=str, default="./outputs")
    parser.add_argument('--model_name', type=str, default="resnet18", choices=['resnet18', 'efficientnet_b0'])
    parser.add_argument('--optimizer', type=str, default="adam", choices=['sgd', 'adam', 'radam'])
    parser.add_argument('--lr', type=float, default=1e-3)
    parser.add_argument('--scheduler', action='store_true')
    parser.add_argument('--batch_size', type=int, default=32)
    parser.add_argument('--img_size', type=int, default=224)
    parser.add_argument('--epochs', type=int, default=10)
    parser.add_argument('--threshold', type=float, default=0.5)

    args = parser.parse_args()
    train_model(args)