AI实战课：利用TensorFlow预测你能否月薪过万

我正在参加「创意开发 投稿大赛」详情请看：掘金创意开发大赛来了！

阅读本文，你将收获以下内容：

1、看到一份2万人的数据，包含年龄、学历、性别、职业、婚姻状况，以及是否年薪过万。

2、使用TensorFlow框架，结构化数据，建立模型并学习训练，全部代码将不超过100行。

3、输入自己的情况，查看预测结果，调整部分特征数据，查看哪些参数对结果影响较大。

序

月薪过万这个词很有意思。有的人觉得很难做到，有的人觉得很容易做到。正所谓：会者不难，难者不会。

那么，月薪过万究竟和什么有关系呢？学历？年龄？性别？

我这里有一份大约2万人的数据。

这类数据不难找，专门做数据研究的人都知道，好多官方的机构（国外居多）都有公开的数据集，供民众研究和学习使用。就像下面这样，各行各业的都有。

那么，看我手里的这份数据，如果让我来分析，我肯定是这个思路，首先是否月薪过万，和学历有关系，博士肯定都月薪过万。

但是，从数据来看，确实存在很多月薪不过万的博士。

继续观察发现，这些博士虽然学历高，但是多数是个体户，由此可见干个体的很难月薪过万。

但是看下面数据，很多个体户也能月薪过万。

我又发现上面结果中出生地是城市的居多，那是不是从城市出生的，从小就受到了城市化的影响。所以，城市出生的干个体的更容易月薪过万呢？

我这么研究，肯定会陷入无限分裂之中，这只是8个字段，如果是那种预测癌症上百个字段的，人工肯定是无法完成的。

那我们不如把他交给人工智能来处理，让AI去学习研究和推断。

一、读入数据

以下代码全部基于python 3.8 + tensorflow 2.3，运行如有报错，请注意查看版本是否对应。

首先，第一步是导入相关的包，看导入的模块就知道，这是一个近似HelloWorld级别的程序。所以大家不用恐惧，所有代码真的不会超过100行（超过了，你来评论区喷我）。

import tensorflow as tf
from tensorflow import feature_column
from tensorflow.keras import layers

from sklearn.model_selection import train_test_split
import pandas as pd
import os

除了TensorFlow外，注意安装pandas和sklearn。

把我们的test.csv数据集放到代码的同级目录下，然后调用下面代码就可以把数据读入到内存了。

csv_file = 'test.csv'
dataframe = pd.read_csv(csv_file)
dataframe

如果你也是用的vscode，并且也用jupyter调试，运行上面的代码可以输出如下表格。

这里我要做一个特殊说明：我把中文的字符都改成了字母，因为程序确实很难更好地支持中文（编码和字节占用问题）。

1.1 命名之争

我没有说把中文改成英文，而是改成了字母，因为有些翻译用的是汉语拼音（先别喷）。在中英文的命名方面，我不认为纯英文就是高大上，定义一个拼音变量就是low。我的方法论就是：谁能做到既简洁又明确，我就用谁。

比如上面的单位性质一栏分为：国企，私企，个体户。我们来看一下英文和拼音的对照。

名称	英文	拼音
国企	state enterprise	guoqi
私企	private enterprise	siqi
个体户	individual business	geti

首先看长度，拼音要明显短得多（你是否也讨厌代码里一半字符都是变量的命名，官方代码这样是因为他们想不到更好的方法，缩写会影响可读性），就算英文缩写也无法做到这么短。其次看表意，在单位性质一栏，如果我说guoqi大家肯定想到的是私企国企里的国企，而非“过期”、“国戚”之类的，所以表意也明确。那为什么不选拼音呢？

其实，对于计算机来说，最好处理的是数字。

我们喜欢的	代码喜欢的	计算机喜欢的
男	man	1
女	woman	0

但是，这样也会有一个问题，如果只有男、女我们还记得住，如果是职业，有几十种之多，来一个13，请问这是什么职业，你总不能拿出字典来查吧，所以这里面免不了各种转化。

二、组织数据集

数据读入后，就需要组装成TensorFlow框架需要的格式了。

# 所有数据，20%化为测试集，80%是训练集
train, test = train_test_split(dataframe, test_size=0.2)
# 训练集里面，再分20%是验证集，80%是训练集
train, val = train_test_split(train, test_size=0.2)

# 将数据组合成(输入,输入)2项
def df_to_dataset(dataframe, shuffle=True, batch_size=32):
    dataframe = dataframe.copy()
    labels = dataframe.pop('result')
    ds = tf.data.Dataset.from_tensor_slices((dict(dataframe), labels))
    if shuffle:
        ds = ds.shuffle(buffer_size=len(dataframe))
    ds = ds.batch(batch_size)
    return ds
    
# 每32组数据为一个批次，将3类数据集都做处理
batch_size = 32
train_ds = df_to_dataset(train, batch_size=batch_size)
val_ds = df_to_dataset(val, shuffle=False, batch_size=batch_size)
test_ds = df_to_dataset(test, shuffle=False, batch_size=batch_size)

这里面需要讲的，也是大家感兴趣的，那就是输入和输入到底是怎么组装的。

输入是8个字段（年龄、单位性质、学历、婚姻状况、职业、性别、一周工作小时数、出生地），输出为是否月薪过万（是或者否）。

关键代码就是把读入的数据，先调用labels = dataframe.pop('result')把结果列分割出来存到labels标签里面。然后，把剩余的其他字段通过dict(dataframe)进行字典化。

我们来跟踪一下这两句核心代码，看看从数据层面它发生了什么，来，上jupyter。

相信不用我多说什么了，使用jupyter调试就是这么强大，每一行的代码你都能看到它执行了什么。最终就是把数据组成了字典（key是字符串，value是对象）的输入+数字结果（1或0）的输出。

三、构建模型的神经网络

我们再来看一下我们的数据集。

输入项总共有8列，我们也已经把它搞成了字典。

但是，如何把它们交给神经网络的处理层，这是一门学问，需要看我们的设计和构思。

我们首先要定义一个特征列的数组，然后把我们需要关注的列加进去，然后构建一个特征层，把这个特征层放入神经网络序列中，然后它们才能运转。

feature_columns = [] # 存放特征列
feature_columns.append(col_a) # A列的特征
feature_columns.append(col_b) # B列的特征
……
# 构建特征层
feature_layer = tf.keras.layers.DenseFeatures(feature_columns)
model = tf.keras.Sequential([feature_layer,……]) # 模型
model.compile(……) # 配置
model.fit(……) # 开始训练

3.1 数值列

比如，每周工作小时数hours这一列，我就想看看具体的数值对结果的影响，那么我们就把它定义成数值列。

# 将hours设置为数值列
hours = feature_column.numeric_column("hours")
feature_columns.append(hours)

3.2 分桶列

但是，有时候我们并不关注具体的数值，比如年龄，数据中从17岁到90岁都有。我们想了解某个年龄段对是否月薪过万的影响，这时候就适合用分桶列。

age = feature_column.numeric_column("age")
age_buckets = feature_column.bucketized_column(age, boundaries=[20,25,30,35,40,45,60,70])
feature_columns.append(age_buckets)

其中boundaries参数中的值[20,25,30,35,40,45,60,70]就是把数据中的年龄分成多个桶。

20以下	20到25	25到30	……	70以上
李子涵16岁、刘梓含19岁	肖雨轩24岁	王大锤28岁	……	丁文元89岁（虚岁）

只有在不同桶里才有差别，比如子涵和梓含虽然差3岁，但是我们让神经网络认为他们是一样的。其实这种分类很有用，因为20岁和30岁对于收入可能会有很大差别，但是70岁和90岁可能不会有太明显的差别。

3.3 分类列

即便是我们已经把私企转化为了siqi。但是，计算机还是无法更好地计算。他们真的只喜欢数值。所以，需要把字母再进一步转化，转为one_hot也就是独热形式。

unit = feature_column.categorical_column_with_vocabulary_list('unit', ['guoqi', 'siqi', 'geti'])
unit_one_hot = feature_column.indicator_column(unit)
feature_columns.append(unit_one_hot)

上面代码其实就是做如下处理。

名称	guoqi	siqi	geti
索引	0	1	2
独热表示	[1, 0, 0]	[0, 1, 0]	[0, 0, 1]

根据上面说的，我们想看看这8项输入对结果都有什么影响，最终构建特征层是这样的：

feature_columns = [] # 特征列
# 年龄
age = feature_column.numeric_column("age")
age_buckets = feature_column.bucketized_column(age, boundaries=[20,25,30,35,40,45,60,70])
feature_columns.append(age_buckets)
# 单位性质
unit = feature_column.categorical_column_with_vocabulary_list('unit', ['guoqi', 'siqi', 'geti'])
unit_one_hot = feature_column.indicator_column(unit)
feature_columns.append(unit_one_hot)
# 学历
xueli = feature_column.categorical_column_with_vocabulary_list('xueli', ['gaozhong', 'zhuanke', 'benke', 'shuoshi', 'boshi'])
xueli_one_hot = feature_column.indicator_column(xueli)
feature_columns.append(xueli_one_hot)
# 婚姻
hunyin = feature_column.categorical_column_with_vocabulary_list('hunyin', ['weihun', 'yihun', 'lihun', 'sang_ou'])
hunyin_one_hot = feature_column.indicator_column(hunyin)
feature_columns.append(hunyin_one_hot)
# 职业
zhiye = feature_column.categorical_column_with_vocabulary_list('zhiye', ['guanli', 'jiaoxue', 'jishu', 'nongmin', 'sale', 'siji', 'weixiu', 'wenyuan', 'wuye'])
zhiye_one_hot = feature_column.indicator_column(zhiye)
feature_columns.append(zhiye_one_hot)
# 性别
sex = feature_column.categorical_column_with_vocabulary_list('sex', ['man', 'woman'])
sex_one_hot = feature_column.indicator_column(sex)
feature_columns.append(sex_one_hot)
# 一周工作时长
feature_columns.append(feature_column.numeric_column("hours"))
# 出生地
address = feature_column.categorical_column_with_vocabulary_list('address', ['village', 'city'])
address_one_hot = feature_column.indicator_column(address)
feature_columns.append(address_one_hot)
# 特征层
feature_layer = tf.keras.layers.DenseFeatures(feature_columns)

四、开始训练

其实，最难的（数据准备和处理）我们都做完了。

训练反而是最简单的，只需要构建好模型，然后配置一下，训练就可以了。

# 构建模型
model = tf.keras.Sequential([
  feature_layer,
  layers.Dense(64, activation='relu'),
  layers.Dense(32, activation='relu'),
  layers.Dense(1, activation='sigmoid')
])
# 配置
model.compile(optimizer='adam',loss='binary_crossentropy',metrics=['accuracy'], run_eagerly=True)
# 存放训练结果
model.load_weights('model.ckpt')  
cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath='model.ckpt', save_best_only=True)
# 进行训练
model.fit(train_ds,validation_data=val_ds,epochs=100, callbacks=[cp_callback])

我们选择使用relu激活函数，第一层64个神经元，第二层32个，最后一层1个输出（和训练数据的标签维度对应），输出采用的是sigmoid激活函数，它将输出一个0到1之间的数，我们就当它是月薪过万可能性的百分比吧。

训练完毕之后，同级目录下应该有model.ckpt系列文件，那就是训练的结果。

因为时间有限，我只训练了100轮。最后，我试了一下测试集的效果。

loss, accuracy = model.evaluate(test_ds)
print("Accuracy", accuracy)
# 120/120 [==============================] - 9s 76ms/step
# loss: 0.4170 - accuracy: 0.7953

预测的准确率已经接近80%了。

也就是说当它遇到一组从来没有见过的数据，它作出预测之后，和标准答案一对比，准确率是80%。

五、验证效果

我拿我的数据试一下。

sample = {
    'age': 31,
    'unit': 'siqi',
    'xueli': 'zhuanke',
    'hunyin': 'yihun',
    'zhiye': 'jishu',
    'sex': 'man',
    'hours': 42,
    'address': 'village'
}
input_dict = {name: tf.convert_to_tensor([value]) for name, value in sample.items()}
predictions = model.predict(input_dict)
prob = tf.nn.sigmoid(predictions[0])

print("%.1f ok." % (100 * prob)) # 60.2 ok.

结果显示我有60%可能性月薪过万。

我再改改参数试试。

修改内容	结果变化	说明
原始数据	60.2%	31岁已婚在私企不加班的专科IT农村男
年龄设为50	63.9%	年龄变大更有可能涨薪
出生地改为城市	60.8%	IT行业和出生在哪里关系不大
婚姻改为未婚	52.2%	程序员不结婚不容易月薪过万
工作单位改为国企	59.9%	去国企有降低工资的风险

从大数据看，我还得好好学习，等着年龄增长了，起码秃顶了，才能挣到更多的钱。

结

注意：本例子并不是最佳实践，因为如此少量的字段和数据量，有更简单和有效的算法来实现，比如决策树或者随机森林。本文所述方法适合做为更高级和复杂数据的起点，此处只是从教学的角度来讲述如何对样本数据进行结构化处理。

完整代码和数据集文件都已开源到GitHub：https://github.com/hlwgy/pay_forecast,欢迎大家来学习和交流。