NCU-机器学习-作业4:基于XGboost的收入分类预测

时间:2024-10-18 08:25:00

输入数据:

在train/目录下包含一个train.csv文件,其中每行表示数据的一条记录。文件中每行共有18列,列income为收入分类(0表示小于等于50k,1表示大于50k)。其余各列为数据对象的各项参数特征(连续属性:age、fnlwgt、education-num、capital-gain;workclass、capital-loss、hours-per-week;离散属性:education、marital-status、occupation、relationship、race、sex、native-country)。训练数据集可在教学资料中下载,微文件名:xgboost-train.csv。

在test/目录下包含一个test.csv文件,与train.csv类似,只不过test.csv不包含income列,您需要根据参数给出预测。

输出数据:

您的程序需要生成一个result.csv文件,用于保存您程序的预测结果。第一行固定为income,之后每一行为一个1或0的值表示预测结果。

评价指标:

准确率。

输入样例:

age,workplace,id,education,education_num,marital_status,occupation,relationship,race,sex,capital_gain,capital_loss,hours_per_week,native_country,income
39, State-gov,77516, Bachelors,13, Never-married, Adm-clerical, Not-in-family, White, Male,2174,0,40, United-States,0
50, Self-emp-not-inc,83311, Bachelors,13, Married-civ-spouse, Exec-managerial, Husband, White, Male,0,0,13, United-States,0
38, Private,215646, HS-grad,9, Divorced, Handlers-cleaners, Not-in-family, White, Male,0,0,40, United-States,0

输出样例:

income
1
0
0
0
1

代码:

import pandas as pd
import xgboost as xgb

# 加载数据
train_data = pd.read_csv('train/train.csv')
test_data = pd.read_csv('test/test.csv')

# 分离特征和标签
X_train = train_data.drop('income', axis=1)
y_train = train_data['income']
X_test = test_data
# print(X_test)

# 独热编码类别型特征
X_train = pd.get_dummies(X_train)
X_test = pd.get_dummies(X_test)

# 确保训练集和测试集有相同的列
X_test = X_test.reindex(columns=X_train.columns, fill_value=0)

# 训练XGBoost模型
xgb_model = xgb.XGBClassifier(objective='binary:logistic')
xgb_model.fit(X_train, y_train)

# 预测测试数据
predictions = xgb_model.predict(X_test)

# 将预测结果写入CSV文件
result_df = pd.DataFrame({'income': predictions})
result_df.to_csv('result.csv', index=False)

Tips:仅为样例代码。