tiny-FWFM/pruned_eval_radio_sig_identification.py at main · Mohammad-Hallaq/tiny-FWFM · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
import os


def evaluate_model(model_path, save_path=None, title='Finetuning Results', device=None):
    from tqdm import tqdm
    from torch.utils.data import DataLoader
    from dataset_classes.radio_sig import RadioSignal
    import torch
    import models_vit
    import models_mae
    import numpy as np
    import random
    from pathlib import Path
    import matplotlib.pyplot as plt
    import seaborn as sns
    import timm
    import torch.nn.functional as F
    from sklearn.metrics import confusion_matrix

    plt.rcParams['font.family'] = 'serif'

    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)
        return out.mean(dim=1)

    def forward(self, x):
        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 = self.attn_drop(attn.softmax(dim=-1))
            x = attn @ v

        x = x.transpose(1, 2).reshape(B, N, -1)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

    seed = 42
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)

    dataset_train = RadioSignal(Path('fine-tuning_datasets/radio_sig_identification/train'))
    dataset_test = RadioSignal(Path('fine-tuning_datasets/radio_sig_identification/test'))

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

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

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

    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, type(m))

    model.eval()

    class_names = ['ais', 'bluetooth', 'cellular', 'fm', 'lora',
                   'packet','rke', 'RS41-Radiosonde', 'sstv', 'wifi']

    def get_preds(loader):
        all_labels, all_preds = [], []
        with torch.no_grad():
            for samples, targets in tqdm(loader):
                samples, targets = samples.to(device), targets.to(device)
                output = model(samples)
                all_preds.extend(output.argmax(dim=-1).cpu().numpy())
                all_labels.extend(targets.cpu().numpy())
        return np.array(all_labels), np.array(all_preds)

    labels_train, preds_train = get_preds(dataloader_train)
    labels_test, preds_test = get_preds(dataloader_test)

    def normalize_conf_mat(cm):
        cm = cm.astype(np.float32)
        for i in range(len(cm)):
            cm[i] /= np.sum(cm[i]) if np.sum(cm[i]) > 0 else 1
        return cm

    cm_train = normalize_conf_mat(confusion_matrix(labels_train, preds_train))
    cm_test = normalize_conf_mat(confusion_matrix(labels_test, preds_test))

    acc_train = np.mean(np.diagonal(cm_train))
    acc_test = np.mean(np.diagonal(cm_test))

    fig, axs = plt.subplots(1, 2, figsize=(10, 5))
    fig.suptitle(title)
    sns.heatmap(cm_train, cmap='Reds', yticklabels=class_names, ax=axs[0])
    sns.heatmap(cm_test, cmap='Reds', yticklabels=class_names, ax=axs[1])
    axs[0].tick_params(axis='y', labelsize=10)
    axs[1].tick_params(axis='y', labelsize=10)
    axs[0].tick_params(axis='x', labelbottom=False)
    axs[1].tick_params(axis='x', labelbottom=False)
    axs[0].set_title(f'Train - Accuracy: {acc_train:.2f}')
    axs[1].set_title(f'Test - Accuracy: {acc_test:.2f}')
    plt.tight_layout()

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

    return acc_train, acc_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
    )