如何利用给定代码打印出RMSE最低的模型名称？

核心思路 #

将每个模型的名称与对应的RMSE值关联存储（比如用字典），然后通过内置函数找到RMSE最小的条目，提取并打印模型名称。

完整代码示例 #

假设你已经完成了数据划分、模型拟合和RMSE计算，以下是添加获取最优模型名称的代码：

from sklearn.linear_model import LinearRegression, Lasso, Ridge
from sklearn.metrics import mean_squared_error
import numpy as np

# 假设X_train, X_test, y_train, y_test是已准备好的数据集
X_train, X_test = ... 
y_train, y_test = ...

# 初始化并训练模型
lr = LinearRegression()
lr.fit(X_train, y_train)
y_pred_lr = lr.predict(X_test)
rmse_lr = np.sqrt(mean_squared_error(y_test, y_pred_lr))

lasso = Lasso()
lasso.fit(X_train, y_train)
y_pred_lasso = lasso.predict(X_test)
rmse_lasso = np.sqrt(mean_squared_error(y_test, y_pred_lasso))

ridge = Ridge()
ridge.fit(X_train, y_train)
y_pred_ridge = ridge.predict(X_test)
rmse_ridge = np.sqrt(mean_squared_error(y_test, y_pred_ridge))

# 存储模型名称与RMSE的映射
model_rmse_dict = {
    "LinearRegression": rmse_lr,
    "Lasso": rmse_lasso,
    "Ridge": rmse_ridge
}

# 找到RMSE最低的模型名称
best_model_name = min(model_rmse_dict, key=lambda k: model_rmse_dict[k])

# 打印结果
print(f"RMSE最低的模型是: {best_model_name}")

关键代码解释 #

• 字典存储映射：model_rmse_dict 将模型名称字符串作为键，对应的RMSE值作为值，建立一一对应关系。
• min函数找最优：利用min()函数的key参数，传入lambda表达式指定以字典的值（RMSE）作为比较依据，返回对应的键（模型名称）。

如果你想同时打印模型名称和对应的RMSE值，可以这样做：

best_model_info = min(model_rmse_dict.items(), key=lambda x: x[1])
print(f"RMSE最低的模型是: {best_model_info[0]}，对应的RMSE值为: {best_model_info[1]:.4f}")

内容的提问来源于stack exchange，提问作者Gonggong