pruned_federated_eval_positioning.py

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

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')
# device = 'cpu'
# model = models_vit.__dict__['vit_small_patch16'](global_pool='token', num_classes=3, drop_path_rate=0.1, in_chans=4)
# checkpoint = torch.load(Path('/home/ict317-3/Mohammad/mae/output_dir/checkpoint-40.pth'), map_location='cpu', weights_only=False)
# msg = model.load_state_dict(checkpoint['model'], strict=True)
# print(msg)

model = torch.load('/home/ict317-3/Mohammad/mae/positioning_output_dir/fed_avg_best_model.pth', weights_only=False)

print(model)

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()
# distances_train.sort()
# distances_test.sort()

# cdf_train = np.linspace(0, 1, len(dataset_train))
# cdf_test = np.linspace(0, 1, len(dataset_test))
# idx_90_train = np.argmin(np.abs(cdf_train - 0.1))
# idx_90_test = np.argmin(np.abs(cdf_test - 0.1))
#
# plt.rcParams['font.family'] = 'serif'
# fig, axs = plt.subplots(1, 1)
# axs.plot(distances_train, cdf_train, label='train', linewidth=2, color='r')
# axs.plot(distances_test, cdf_test, label='test', linewidth=2, color='b')
# axs.axhline(y=0.1, linewidth=1, linestyle='--', label='90% likely', color='k')
# axs.axvline(x=distances_train[idx_90_train], linewidth=1, linestyle='--', color='r', alpha=0.8)
# axs.axvline(x=distances_test[idx_90_test], linewidth=1, linestyle='--', color='b', alpha=0.8)
# axs.set_xlabel('Positioning Error (m)')
# axs.set_ylabel('CDF')
# axs.legend(loc='lower right')
# plt.tight_layout()
# # plt.savefig('cdf_positioning.png', dpi=300)
# plt.show()

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)

plt.tight_layout()
plt.savefig('CFFM/pruning_models_small_ViT/pruned_ViT_positioning/hist_positioning.png', dpi=300)
plt.show()