pruned_eval_positioning.py

def evaluate_model(model_path, save_path=None, title='Finetuning Results', device=None):
    
    import os
    from tqdm import tqdm

    import models_vit
    from dataset_classes.positioning import Positioning5G
    import torch
    import numpy as np
    import matplotlib.pyplot as plt
    import random
    from pathlib import Path
    from torch.utils.data import random_split, DataLoader

    import torch.nn.functional as F
    import timm

    import models_mae

    def no_mask_forward(self, imgs, mask_ratio=0.0):
            latent, _, ids_restore = self.forward_encoder(imgs, mask_ratio)
            out = self.forward_decoder(latent, ids_restore)

            cls_output = out.mean(dim=1)  # shape: [B, decoder_pred_dim]  
            return cls_output

    def forward(self, x):
        """https://github.com/huggingface/pytorch-image-models/blob/054c763fcaa7d241564439ae05fbe919ed85e614/timm/models/vision_transformer.py#L79"""
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)
        q, k = self.q_norm(q), self.k_norm(k)

        if self.fused_attn:
            x = F.scaled_dot_product_attention(
                q, k, v,
                dropout_p=self.attn_drop.p,
            )
        else:
            q = q * self.scale
            attn = q @ k.transpose(-2, -1)
            attn = attn.softmax(dim=-1)
            attn = self.attn_drop(attn)
            x = attn @ v

        x = x.transpose(1, 2).reshape(B, N, -1) # original implementation: x = x.transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


    def reverse_normalize(x, coord_min, coord_max):
        return (x + 1) / 2 * (coord_max - coord_min) + coord_min


    seed = 42
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    scene = 'outdoor'
    dataset_train = Positioning5G(Path('fine-tuning_datasets/5G_NR_Positioning/outdoor/train'), scene=scene)
    dataset_test = Positioning5G(Path('fine-tuning_datasets/5G_NR_Positioning/outdoor/test'), scene=scene)
    coord_min, coord_max = dataset_train.coord_nominal_min.view((1, -1)), dataset_train.coord_nominal_max.view((1, -1))

    dataloader_train = DataLoader(dataset_train, batch_size=256, shuffle=False, num_workers=0)
    dataloader_test = DataLoader(dataset_test, batch_size=256, shuffle=False, num_workers=0)

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

    model = torch.load(model_path, weights_only=False)

    print(model)

    if model.__class__.__name__ == 'MaskedAutoencoderViT':
                model.forward = no_mask_forward.__get__(model, models_mae.MaskedAutoencoderViT)


    for m in model.modules():
                if isinstance(m, timm.models.vision_transformer.Attention):
                    m.forward = forward.__get__(m, timm.models.vision_transformer.Attention) 


    model = model.to(device)
    distances_train = torch.zeros((len(dataset_train),))

    with torch.no_grad():
        for i, batch in tqdm(enumerate(dataloader_train), desc='Train Batch', total=len(dataloader_train)):
            image, target = batch
            image = image.to(device)
            pred_position = reverse_normalize(model(image).cpu(), coord_min, coord_max)
            position = reverse_normalize(target.cpu(), coord_min, coord_max)
            num_samples = target.shape[0]
            distances_train[i * num_samples: (i + 1) * num_samples] = torch.sqrt(torch.sum((pred_position - position) ** 2, dim=1))


    distances_test = torch.zeros((len(dataset_test),))
    with torch.no_grad():
        for i, batch in tqdm(enumerate(dataloader_test), desc='Test Batch', total=len(dataloader_test)):
            image, target = batch
            image = image.to(device)
            pred_position = reverse_normalize(model(image).cpu(), coord_min, coord_max)
            position = reverse_normalize(target.cpu(), coord_min, coord_max)
            num_samples = target.shape[0]
            distances_test[i * num_samples: (i + 1) * num_samples] = torch.sqrt(torch.sum((pred_position - position) ** 2, dim=1))

    distances_train = distances_train.numpy()
    distances_test = distances_test.numpy()
    

    plt.rcParams['font.family'] = 'serif'
    mean_train = np.mean(distances_train)
    mean_test = np.mean(distances_test)

    fig, axs = plt.subplots(1, 2, figsize=(10, 5))
    # model = 'Finetuning ViT-M'
    # other = '(2 out of 12 blocks + linear layer)'
    # fig.suptitle(f'{model} {other}\n{scene} scenario')
    bins = 25
    axs[0].hist(distances_train, bins=bins, color='red', edgecolor='w', alpha=0.7, density=True)
    axs[0].axvline(mean_train, color='black', linestyle='--', linewidth=2, label=f'Mean: {mean_train:.2f} (m)')
    # axs[0].set_title('Training')
    axs[0].set_xlabel('Positioning Error (m)', fontsize=16)
    axs[0].set_ylabel('Probability Density', fontsize=16)
    axs[0].legend(fontsize=16)

    axs[1].hist(distances_test, bins=bins, color='blue', edgecolor='w', alpha=0.7, density=True)
    axs[1].axvline(mean_test, color='black', linestyle='--', linewidth=2, label=f'Mean: {mean_test:.2f} (m)')
    # axs[1].set_title('Test')
    axs[1].set_xlabel('Positioning Error (m)', fontsize=16)
    axs[1].set_ylabel('Probability Density', fontsize=16)
    axs[1].legend(fontsize=16)

    if save_path:
        plt.savefig(os.path.join(save_path), dpi=300)
    plt.show()

    return mean_train, mean_test

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--model_path', type=str, required=True, help='Path to the trained model (.pth)')
    parser.add_argument('--save_path', type=str, default=None, help='Optional path to save the confusion matrix image')
    parser.add_argument('--title', type=str, default='Finetuning Results', help='Title for the confusion matrix plot')

    args = parser.parse_args()

    evaluate_model(
        model_path=args.model_path,
        save_path=args.save_path,
        title=args.title
    )