model_utils.py

import torch
from torch import nn
from torch.optim import *
from torch.optim.lr_scheduler import *
from torch.utils.data import DataLoader
from torchvision.datasets import *
from torchvision.transforms import *
from typing import Union, Any, Tuple
import copy
import torch.nn.utils.prune as prune
import logging
import psutil, os, gc
import typing, functools
from collections.abc import Iterator

from torchvision.models import resnet50, ResNet50_Weights
from torchvision.models import resnet34, ResNet34_Weights
from torchvision.models import resnet18, ResNet18_Weights
from torchvision.models import vgg19_bn, VGG19_BN_Weights   #.IMAGENET1K_V1


def load_model(model_name: str, 
               dataset_name: str, 
               device: Union[torch.device, str] = torch.device('cuda'),
               ) -> nn.Module:
    
    if 'cifar' in dataset_name:
        full_name = dataset_name + "_" + model_name
        model = torch.hub.load("chenyaofo/pytorch-cifar-models", full_name, pretrained=True)

    elif 'imagenet' in dataset_name:
        if model_name == "vgg19_bn":
            model = vgg19_bn(weights=VGG19_BN_Weights.DEFAULT).to(device)

        elif model_name == "resnet50":
            model = resnet50(weights=ResNet50_Weights.DEFAULT).to(device)

        elif model_name == "resnet34":
            model = resnet34(weights=ResNet34_Weights.DEFAULT).to(device)
        
        elif model_name == "resnet18":
            model = resnet18(weights=ResNet18_Weights.DEFAULT).to(device)

    return model.to(device)


def load_params(model: nn.Module, 
                addr: str, 
                device: Union[torch.device, str] = torch.device('cpu'),
                ) -> None:
    i = 0
    state_dict_load = torch.load(addr, map_location=device)
    sd = model.state_dict()
    for layer_name, _ in sd.items():
        if 'num_batches_tracked' not in layer_name:
            sd[layer_name] = nn.Parameter(state_dict_load[list(state_dict_load)[i]].to(device))
        i += 1
    
    model.load_state_dict(sd)

def load_params_sparse(model: nn.Module, 
                addr: str, 
                device: Union[torch.device, str] = torch.device('cpu'),
                ) -> None:
    i = 0
    state_dict_load = torch.load(addr, map_location=device)
    sd = model.state_dict()
    for layer_name, _ in sd.items():
        if 'num_batches_tracked' not in layer_name:
            state_dict_name = layer_name
            if "weight" in layer_name and layer_name not in state_dict_load.keys():
                state_dict_name += "_orig"
            sd[layer_name] = nn.Parameter(state_dict_load[state_dict_name].to(device))
    
    model.load_state_dict(sd)

def train(model: nn.Module,
          dataloader: DataLoader,
          criterion: nn.Module,
          optimizer: Optimizer,
          scheduler: LambdaLR,
          callbacks = None,
          device=torch.device('cuda')) -> None:
    
    model.train()

    for data in dataloader:
        inputs, targets = data[0].to(device), data[1].to(device)
        optimizer.zero_grad()

        # Forward inference
        outputs = model(inputs)
        loss = criterion(outputs, targets)

        # Backward propagation
        loss.backward()

        # Update optimizer and LR scheduler
        optimizer.step()
        scheduler.step()

        if callbacks is not None:
            for callback in callbacks:
                callback()

@functools.cache
def cache_data(dataloader, device):
    batches: list[tuple[torch.Tensor, torch.Tensor]] = []
    iterator = typing.cast(Iterator[list[torch.Tensor]], iter(dataloader))
    for batch in iterator:
        assert isinstance(batch, list)
        assert len(batch) == 2
        assert all(isinstance(x, torch.Tensor) for x in batch)

        batches.append(
            (batch[0].to(torch.float).to(device), batch[1].to(torch.float).to(device))
        )
    
    return batches

#@torch.inference_mode()
def evaluate(
    model: nn.Module,
    dataloader,
    device=torch.device('cuda')) -> float:
    
    model.to(device)
    model.eval()
    num_samples = 0
    num_correct = 0

    with torch.no_grad():
        for data in dataloader:
            inputs, targets = data[0].to(device), data[1].to(device)

            # Inference
            outputs = model(inputs)

            # Convert logits to class indices
            results = outputs.argmax(dim=1)

            # Update metrics
            num_samples = num_samples + targets.size(0)
            num_correct = num_correct + (results == targets).sum().item()

            del outputs
            del results
            torch.cuda.empty_cache()

    return (num_correct / num_samples * 100)

def faulty_evaluate(
    golden_model: nn.Module,
    faulty_model: nn.Module,
    dataloader,
    device=torch.device('cuda'),
    classes_count: int=10):
    
    golden_model.to(device)
    faulty_model.to(device)
    golden_model.eval()
    faulty_model.eval()
    counter = 0
    SDC, SDC_critical, faulty_accuracy = 0, 0, 0
    
    with torch.no_grad():
        for data in dataloader:
            inputs, targets = data[0].to(device), data[1].to(device)

            # Inference
            golden_outputs = golden_model(inputs)
            faulty_outputs = faulty_model(inputs)

            # Convert logits to class indices
            golden_results = golden_outputs.argmax(dim=1)
            faulty_results = faulty_outputs.argmax(dim=1)

            faulty_results = torch.nan_to_num(faulty_results, nan=classes_count+1, posinf=classes_count+1, neginf=classes_count+1)
            
            # Update metrics
            # num_samples += targets.size(0)
            faulty_accuracy += (faulty_results == targets).sum().item() / (targets.size(0))
            SDC += torch.sum(faulty_outputs != golden_outputs).item() / (torch.numel(faulty_outputs))
            SDC_critical += torch.sum(faulty_results != golden_results).item() / (targets.size(0))
            counter += 1

            del faulty_results
            del faulty_outputs
            del golden_outputs

            torch.cuda.empty_cache()

    return faulty_accuracy * 100 / counter, SDC * 100 / counter, SDC_critical * 100 / counter

def evaluate_partial_batch(
    model: nn.Module,
    dataloader,
    device=torch.device('cuda')) -> float:
    
    model.to(device)
    model.eval()
    num_samples = 0
    num_correct = 0
    batch_count = 0
    max_batch = 10

    with torch.no_grad():
        for data in dataloader:
            inputs, targets = data[0].to(device), data[1].to(device)

            # Inference
            outputs = model(inputs)

            # Convert logits to class indices
            results = outputs.argmax(dim=1)

            # Update metrics
            num_samples += targets.size(0)
            num_correct += (results == targets).sum()
            batch_count += 1

            del outputs
            del results
            torch.cuda.empty_cache()

            if batch_count < max_batch:
                break


    return (num_correct / num_samples * 100).item()

def fine_tune_sparse(model: nn.Module,
              trainloader: DataLoader,
              testloader: DataLoader,
              eopchs: int,
              lr: float,
              device: Union[torch.device, str],
              saved_model_name: str,
              logger: logging.Logger
              ) -> nn.Module:
    
    optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, eopchs)
    criterion = nn.CrossEntropyLoss()

    best_accuracy = 0
    # best_epoch = -1
    for epoch in range(eopchs):
        train(model, trainloader, criterion, optimizer, scheduler, device=device)
        accuracy = evaluate_partial_batch(model, testloader, device=device)
        if accuracy > best_accuracy:
            best_accuracy = accuracy
            sparse_model = copy.deepcopy(model)
            for _, module in sparse_model.named_modules():
                if isinstance(module, torch.nn.Conv2d):
                    if hasattr(module, 'weight_mask'):
                        prune.remove(module, 'weight')
            torch.save(sparse_model.state_dict(), saved_model_name)
            best_accuracy = accuracy
            #best_epoch = epoch
            logger.info(f"best accuracy until epoch {epoch}: {accuracy}")
            
            del sparse_model
            torch.cuda.empty_cache()
    #     logger.info(f"epoch {epoch}, accuracy: {accuracy}%")
    
    # accuracy = evaluate(model, testloader, device=device)
    # logger.info(f"model saved with the accuracy {accuracy}%")

    return model

def memory_stats(i):
    print(f"GPU memory {i}: {torch.cuda.memory_allocated()/1024**2}")
    print(f"GPU cache {i}: {torch.cuda.memory_cached()/1024**2}")
    process = psutil.Process(os.getpid())
    print(f"CPU memory {i}: {process.memory_info().rss / 1024**2}")
    print(f"Uncollected objects: {len(gc.garbage)}")
    print("--------------")