eval_memory.py

# Third-party imports
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import v2 as transforms
from transformers import MobileViTV2ForImageClassification, MobileViTV2Config
from timm import create_model
from GMR_Conv import gmr_resnet18

# Local imports
from model_hybrid import HybridCNNViT
from SingleModels.e2cnn_model import E2WRN16_8_Star_D8D4D1_Cls
from utils import load_config
from cli_args import get_opt

args = get_opt()

if torch.cuda.is_available():
    device = "cuda"
elif torch.backends.mps.is_available():
    device = "mps"
else:
    device = "cpu"

cfg = load_config("config/train.yaml")
IMG_SIZE = args.img_size

transform = transforms.Compose([
    transforms.ToImage(),
    transforms.Resize(IMG_SIZE),
    transforms.ToDtype(torch.float32, scale=True),
    transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]),
])

test_data = datasets.CIFAR10(
    root=cfg["SAVE_PATH"]["DATA"],
    train=False,
    download=True,
    transform=transform
)
test_loader = DataLoader(test_data, batch_size=cfg["BATCH_SIZE"]["TEST"])

def build_model(model_name, ckpt):
    if model_name == "gmr":
        model = gmr_resnet18(
            num_classes=cfg["NUM_CLASSES"],
            inplanes=64,
            in_channels=3,
            gmr_conv_size=[9, 9, 5, 5],
            num_rings=None,
        )
        # for pcam
        model.conv1 = nn.Conv2d(
            3,
            64,
            kernel_size=5,
            stride=1,
            padding=2,
            bias=False,
        )
        model = model.to(device)
        # model.load_state_dict(torch.load("models/gmr.pt", weights_only=True))
        checkpoint = torch.load("models/gmr_pcam_no_rot.ckpt", map_location=device)
        model.load_state_dict(checkpoint["model"])
        return model.eval()

    if model_name == "e2cnn":
        model = E2WRN16_8_Star_D8D4D1_Cls(
            num_classes=cfg["NUM_CLASSES"],
        ).to(device)
        # model.load_state_dict(torch.load("models/e2cnn.pt", weights_only=True))
        checkpoint = torch.load("models/e2cnn_pcam.ckpt", map_location=device)
        model.load_state_dict(checkpoint["model"])
        return model.eval()

    if model_name == "mobile":
        hf_model = "models/mobilevitv2-1.0"
        config = MobileViTV2Config.from_pretrained(hf_model, local_files_only=True)
        config.num_labels = cfg["NUM_CLASSES"]
        model = MobileViTV2ForImageClassification(config).to(device)

        model.load_state_dict(torch.load("models/mobile.pt", weights_only=True))
        return model.eval()

    if model_name == "timm":
        ckpt = "models/efficientvit_m2_pcam_no_rot.pt" #"models/shvits1.pt" #"models/efficientvit_m2.pt"
        timm_model = "efficientvit_m2" #"shvit_s1" #"efficientvit_m2"
        model = create_model(timm_model, pretrained=False, num_classes=cfg["NUM_CLASSES"]).to(device)
        model.load_state_dict(torch.load(ckpt, weights_only=True))
        return model.eval()

    if model_name == "hybrid":
        model = HybridCNNViT(num_classes=cfg["NUM_CLASSES"]).to(device)

        model.load_state_dict(
            torch.load(ckpt, weights_only=True))  # cfg["SAVE_PATH"]["MODEL"]
        # checkpoint = torch.load(cfg["SAVE_PATH"]["CHECKPOINTS"], map_location=device)
        # model.load_state_dict(checkpoint["model"])
        return model.eval()

    raise ValueError(f"Unknown model: {model_name}")


torch.manual_seed(cfg["SEED"])
model = build_model(args.model, args.ckpt)

dummy = torch.randn(cfg["BATCH_SIZE"]["TEST"], 3, IMG_SIZE, IMG_SIZE).to(device)
warmup, iters = 30, 200
with torch.inference_mode():
    # Warmup: trigger kernel selection / lazy init
    for _ in range(warmup):
        _ = model(dummy) #.logits
    torch.cuda.synchronize()

    # Start measurement window
    torch.cuda.reset_peak_memory_stats(device)

    for _ in range(iters):
        _ = model(dummy) #.logits
    torch.cuda.synchronize()

peak_alloc_mib = torch.cuda.max_memory_allocated(device) / (1024 ** 2)
peak_resv_mib  = torch.cuda.max_memory_reserved(device) / (1024 ** 2)

#print(f"Peak allocated (MB): {peak_alloc_mib:.2f} MiB")
print(f"Peak reserved  (MB): {peak_resv_mib:.2f} MiB")

total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

print("Total parameters:", total_params)
print("Trainable parameters:", trainable_params)