quantization-of-tgn/generate_plots.py at main · an-andrey/quantization-of-tgn · 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
import os
import argparse
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

def generate_plots(results_dir, dataset_name):
    csv_path = os.path.join(results_dir, "all_results_merged.csv")

    if not os.path.exists(csv_path):
        print(f"Error: CSV file not found at {csv_path}")
        return

    df = pd.read_csv(csv_path)
    df_dataset = df[df['Dataset'] == dataset_name]

    if df_dataset.empty:
        print(f"No data found for dataset: {dataset_name}")
        return

    plots_dir = os.path.join(results_dir, "plots")
    os.makedirs(plots_dir, exist_ok=True)

    # Styling parameters
    plt.rcParams.update({'font.size': 12, 'font.family': 'serif'})
    colors = {'fp32': '#1f77b4', 'fp16': '#ff7f0e', 'int8': '#2ca02c'}
    precisions = df_dataset['Precision'].unique()

    # ---------------------------------------------------------
    # Plot 1: Test MRR vs Epoch (Line Plot)
    # ---------------------------------------------------------
    fig1, ax1 = plt.subplots(figsize=(8, 5))

    for prec in precisions:
        subset = df_dataset[df_dataset['Precision'] == prec].sort_values(by='Epoch')
        ax1.plot(subset['Epoch'], subset['Test_MRR'], marker='o', label=prec.upper(), color=colors.get(prec, '#333333'))

    ax1.set_xlabel('Epoch')
    ax1.set_ylabel('Test MRR')
    ax1.set_title(f'Test MRR Convergence: {dataset_name}')
    ax1.grid(True, linestyle='--', alpha=0.7)
    ax1.legend()

    mrr_plot_path = os.path.join(plots_dir, f"{dataset_name}_mrr_convergence.png")
    fig1.tight_layout()
    fig1.savefig(mrr_plot_path, dpi=300)
    plt.close(fig1)
    print(f"Saved MRR plot to {mrr_plot_path}")

    # ---------------------------------------------------------
    # Plot 2: Payload and Peak CPU RAM (Grouped Bar Chart)
    # ---------------------------------------------------------
    fig2, ax2 = plt.subplots(figsize=(8, 6))
    ax3 = ax2.twinx()

    x = np.arange(len(precisions))
    width = 0.35

    payloads = []
    cpu_rams = []

    for prec in precisions:
        subset = df_dataset[df_dataset['Precision'] == prec]
        # Memory metrics remain static across epochs, take the max or mean
        payloads.append(subset['Total_Payload_MB'].max())
        cpu_rams.append(subset['Peak_CPU_MB'].max())

    bars1 = ax2.bar(x - width/2, payloads, width, label='Edge Feature Payload', color='#4c72b0')
    bars2 = ax3.bar(x + width/2, cpu_rams, width, label='Peak CPU RAM', color='#c44e52')

    ax2.set_xlabel('Precision State')
    ax2.set_ylabel('Edge Feature Payload (MB)', color='#4c72b0')
    ax3.set_ylabel('Peak CPU RAM (MB)', color='#c44e52')
    ax2.set_title(f'Spatial Memory Footprint: {dataset_name}')

    ax2.set_xticks(x)
    ax2.set_xticklabels([p.upper() for p in precisions])

    # Combine legends from both axes
    lines_labels = [ax.get_legend_handles_labels() for ax in [ax2, ax3]]
    lines, labels = [sum(lol, []) for lol in zip(*lines_labels)]
    ax2.legend(lines, labels, loc='upper right')

    memory_plot_path = os.path.join(plots_dir, f"{dataset_name}_memory_footprint.png")
    fig2.tight_layout()
    fig2.savefig(memory_plot_path, dpi=300)
    plt.close(fig2)
    print(f"Saved Memory Footprint plot to {memory_plot_path}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate plots for TGN benchmark.")
    parser.add_argument("--results-dir", type=str, required=True, help="Directory containing the merged CSV")
    parser.add_argument("--dataset", type=str, required=True, help="Name of the dataset to plot")

    args = parser.parse_args()
    generate_plots(args.results_dir, args.dataset)