EasyMachineLearningDemo / visualization /draw_learning_curve_total.py
LLH
2024/02/14/01:14
bd39f54
raw
history blame
2 kB
import numpy as np
from matplotlib import pyplot as plt
from static.config import Config
def draw_learning_curve_total(input_dict, type):
plt.figure(figsize=(10, 6), dpi=300)
if type == "train":
i = 0
for label_name, values in input_dict.items():
train_sizes = values[0]
train_scores_mean = values[1]
train_scores_std = values[2]
test_scores_mean = values[3]
test_scores_std = values[4]
plt.fill_between(
train_sizes,
train_scores_mean - train_scores_std,
train_scores_mean + train_scores_std,
alpha=0.1,
color=Config.COLORS[i]
)
plt.plot(
train_sizes,
train_scores_mean,
"o-",
color=Config.COLORS[i],
label=label_name
)
i += 1
title = "Training Learning curve"
# plt.title(title)
else:
i = 0
for label_name, values in input_dict.items():
train_sizes = values[0]
train_scores_mean = values[1]
train_scores_std = values[2]
test_scores_mean = values[3]
test_scores_std = values[4]
plt.fill_between(
train_sizes,
test_scores_mean - test_scores_std,
test_scores_mean + test_scores_std,
alpha=0.1,
color=Config.COLORS[i]
)
plt.plot(
train_sizes,
test_scores_mean,
"o-",
color=Config.COLORS[i],
label=label_name
)
i += 1
title = "Cross-validation Learning curve"
# plt.title(title)
plt.xlabel("Sizes")
plt.ylabel("Adjusted R-square")
plt.legend()
# plt.savefig("./diagram/{}.png".format(title), dpi=300)
# plt.show()
return plt