Dual_I2CL/visualize.py at Lagrangian · Parkprogrammer/Dual_I2CL · 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
import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors

datasets = ['agnews', 'dbpedia', 'emo', 'hate_speech18', 'mr', 'sst2', 'sst5', 'subj', 'trec']
settings = ['zero_shot', 'few_shot', 'i2cl']
base_path = "~/claire/I2CL/fim_exps/i2cl/EleutherAI/gpt-j-6B"

base_path = os.path.expanduser(base_path)

layers = [str(i) for i in range(28)] + ['final_norm']
n_layers = len(layers)

valid_columns = []
valid_labels = []

n_datasets = len(datasets)
n_settings = len(settings)
fim_matrix_full = np.zeros((n_layers, n_datasets * n_settings))

def parse_fim_file(file_path):
    fim_values = {}
    with open(file_path, 'r') as f:
        for line in f:
            if 'Layer:' in line and 'FIM:' in line:
                parts = line.split(',')
                layer = parts[0].split(':')[1].strip()
                fim = float(parts[1].split(':')[1].strip())
                fim_values[layer] = fim
    return fim_values

col_idx = 0
for d_idx, dataset in enumerate(datasets):
    for s_idx, setting in enumerate(settings):

        file_name = f"{dataset}_EleutherAI_gpt-j-6B_{setting}_fim.txt"
        file_path = os.path.join(base_path, dataset, "fim_results", file_name)

        if not os.path.exists(file_path):
            print(f"경고: 파일을 찾을 수 없습니다: {file_path}")
            continue

        try:
            fim_values = parse_fim_file(file_path)
        except Exception as e:
            print(f"파일 파싱 중 오류 발생: {file_path}, 오류: {e}")
            continue

        for l_idx, layer in enumerate(layers):
            fim_matrix_full[l_idx, col_idx] = fim_values.get(layer, 0.0)

        valid_columns.append(col_idx)
        valid_labels.append(f"{dataset[:3]}_{setting}")
        col_idx += 1

if valid_columns:
    fim_matrix = fim_matrix_full[:, valid_columns]
else:
    print("유효한 데이터가 없습니다. 시각화를 생성하지 않습니다.")
    exit()

fim_matrix = np.log10(fim_matrix + 1e-10)  # log(0) 방지를 위해 작은 값 추가
fim_min, fim_max = fim_matrix.min(), fim_matrix.max()
fim_matrix = (fim_matrix - fim_min) / (fim_max - fim_min)  # [0, 1]로 정규화

cmap = mcolors.LinearSegmentedColormap.from_list("custom", ["blue", "red"])

fig, ax = plt.subplots(figsize=(max(12, len(valid_columns) * 0.5), 8))
im = ax.imshow(fim_matrix, cmap=cmap, aspect='auto')

ax.set_yticks(np.arange(n_layers))
ax.set_yticklabels(layers)
ax.set_xticks(np.arange(len(valid_columns)))
ax.set_xticklabels(valid_labels, rotation=45, ha='right')

plt.colorbar(im, ax=ax, label='Normalized Log(FIM)')

dataset_boundaries = []
current_dataset = None
for i, label in enumerate(valid_labels):
    dataset_name = label.split('_')[0]
    if current_dataset != dataset_name:
        if i > 0:
            dataset_boundaries.append(i - 0.5)
        current_dataset = dataset_name

for boundary in dataset_boundaries:
    ax.axvline(x=boundary, color='white', linewidth=2)

ax.set_xlabel('Dataset and Setting')
ax.set_ylabel('Layer')
ax.set_title('FIM Values Across Layers, Datasets, and Settings')

plt.tight_layout()

plt.savefig('fim_heatmap.png', dpi=300, bbox_inches='tight')
print("시각화가 'fim_heatmap.png' 파일로 저장되었습니다.")