文章目录
- 训练简单的线性模型,预测电影评分和票房收入
- 数据集
- 加载数据
- 可视化数据
- 分割数据集
- 训练模型
- 可视化结果
- 应用模型
- 保存和加载模型
训练简单的线性模型,预测电影评分和票房收入
这是一个简单的线性模型,用于预测电影的评分和票房收入。不过却是一个很好的入门案例,可以让你了解如何来解决一个回归问题。
数据集
数据集包含了 10000 条电影的信息,包括评分、票房收入。你可以在这里下载数据集。提取码:1234
。文件格式为 csv
,可以使用 pandas
来读取。
加载数据
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error
import joblib
# 请填写你的数据集路径。
df = pd.read_csv('/movie_rating_box_office_dataset.csv')
df.head()
可视化数据
xDataName = 'Movie_Rating'
yDataName = 'Box_Office_Revenue'
plt.figure(figsize=(10, 6))
plt.scatter(df[xDataName], df[yDataName], alpha=0.5, label='Data points')
plt.title('Scatter Plot of Data')
plt.xlabel(xDataName)
plt.ylabel(yDataName)
plt.legend()
plt.show()
xDataName和yDataName是两个变量,分别代表电影评分和票房收入。我们可以看到,这两个变量之间的关系是线性的.
分割数据集
这里我们将数据集分割成训练集和测试集,其中训练集占 80%,测试集占 20%。
X_train, X_test, Y_train, Y_test = train_test_split(df[['Movie_Rating']], df['Box_Office_Revenue'], test_size=0.2, random_state=0)
(X_train.shape, X_test.shape), (Y_train.shape, Y_test.shape)
训练模型
# 初始化模型
model = LinearRegression()
# 训练模型
model.fit(X_train, Y_train)
# 预测
Y_pred = model.predict(X_test)
# 计算均方误差
mse = mean_squared_error(Y_test, Y_pred)
coef = model.coef_[0]
intercept = model.intercept_
# 打印结果
coef, intercept, mse
输出结果为:
(2.0000851616098485, 5.022264732247818, 1.016650048849064)
其中,coef
和 intercept
分别是线性模型的斜率和截距,mse
是均方误差。
这里我们使用均方误差来评估模型的好坏,均方误差越小,说明模型越好。
可视化结果
plt.figure(figsize=(10, 6))
plt.scatter(X_test, Y_test, alpha=0.5, label='Actual')
plt.scatter(X_test, Y_pred, alpha=0.5, label='Predicted')
plt.plot(X_test, model.predict(X_test), color='red', linewidth=2, label='Fitted line')
plt.title('Actual vs Predicted')
plt.xlabel('X')
plt.ylabel('Y')
plt.legend()
plt.show()
应用模型
# 假设你有一个电影评分为 7.5
new_rating = np.array([[7.5]])
# 使用模型进行预测
predicted_revenue = model.predict(new_rating)
# 输出预测的房价
print(predicted_revenue[0])
保存和加载模型
# 保存模型
joblib.dump(model, 'linear_regression_model.pkl')
# 加载模型
loaded_model = joblib.load('linear_regression_model.pkl')
该模型被保存在 linear_regression_model.pkl
文件中,可以使用 joblib.load
来加载模型。