File size: 1,481 Bytes
a80f6e6 | 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 | import json
import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
import os
import glob
# 自动查找所有子目录下的reward_perplexity_stats.json
stat_files = glob.glob(os.path.join('.', '**', 'reward_perplexity_stats.json'), recursive=True)
for stat_file in stat_files:
with open(stat_file) as f:
stats = json.load(f)
rewards = np.array(stats["step_rewards"])
rewards = rewards - np.mean(rewards)
rewards = rewards / np.std(rewards)
perplexities = np.array(stats["step_perplexities"])
inv_ppl = 1 / perplexities
# inv_ppl = perplexities
# 计算皮尔逊相关系数R
if len(inv_ppl) > 1 and np.std(inv_ppl) > 0 and np.std(rewards) > 0:
R = np.corrcoef(inv_ppl, rewards)[0, 1]
else:
R = float('nan')
# 画散点图
plt.figure(figsize=(8,6))
plt.scatter(inv_ppl, rewards, alpha=0.2)
plt.xlabel("1 / Perplexity")
plt.ylabel("Reward (0/1)")
plt.title(f"Reward vs. Inverse Perplexity\n{stat_file}\nR={R:.3f}")
# 拟合线性关系
X = inv_ppl.reshape(-1, 1)
y = rewards
reg = LinearRegression().fit(X, y)
y_pred = reg.predict(X)
plt.plot(inv_ppl, y_pred, color='red', label='Linear fit')
plt.legend()
plt.tight_layout()
plt.savefig(stat_file.replace('.json', '_reward_vs_invperplexity.png'))
plt.close()
print(f"{stat_file}: Linear fit coef: {reg.coef_}, intercept: {reg.intercept_}, R: {R:.3f}") |