train_data = "./../data/adult/adult.train"
test_data = "./../data/adult/adult.test"
train = pd.read_csv(train_data, sep=",", names=["age", "workclass", "fnlwgt", "education", "education_num", "marital_status", "occupation", "relationship", "race","sex", "capital_gain", "capital_loss", "hours_per_week", "native_country", "label"])
print(train.head(5))

   age          workclass  fnlwgt  ... hours_per_week  native_country   label
0   39          State-gov   77516  ...             40   United-States   <=50K
1   50   Self-emp-not-inc   83311  ...             13   United-States   <=50K
2   38            Private  215646  ...             40   United-States   <=50K
3   53            Private  234721  ...             40   United-States   <=50K
4   28            Private  338409  ...             40            Cuba   <=50K

# 定义基本连续的特征，linear 和 dnn 都会使用到
age = tf.feature_column.numeric_column("age")
education_num = tf.feature_column.numeric_column("education_num")
capital_gain = tf.feature_column.numeric_column("capital_gain")
capital_loss = tf.feature_column.numeric_column("capital_loss")
hours_per_week = tf.feature_column.numeric_column("hours_per_week")

# 定义离散特征
workclass = tf.feature_column.categorical_column_with_vocabulary_list(
    key="workclass",
    vocabulary_list=["Private", "Self-emp-not-inc", "Self-emp-inc", "Federal-gov", "Local-gov", "State-gov",
                     "Without-pay", "Never-worked", "?"]
)
education = tf.feature_column.categorical_column_with_vocabulary_list(
    key="education",
    vocabulary_list=["Bachelors", "Some-college", "11th", "HS-grad", "Prof-school", "Assoc-acdm", "Assoc-voc", "9th",
                     "7th-8th", "12th", "Masters", "1st-4th", "10th", "Doctorate", "5th-6th", "Preschool"]

)
marital_status = tf.feature_column.categorical_column_with_vocabulary_list(
    key="marital_status",
    vocabulary_list=["Married-civ-spouse", "Divorced", "Never-married", "Separated", "Widowed", "Married-spouse-absent",
                     "Married-AF-spouse"]
)
relationship = tf.feature_column.categorical_column_with_vocabulary_list(
    key="relationship",
    vocabulary_list=["Wife", "Own-child", "Husband", "Not-in-family", "Other-relative", "Unmarried"]
)

# 定义Hash特征，展示embedding的使用
occupation = tf.feature_column.categorical_column_with_hash_bucket(
    key="occupation",
    hash_bucket_size=1000
)

# age 特征分桶
age_bucket = tf.feature_column.bucketized_column(
    source_column=age,
    boundaries=[18, 25, 30, 35, 40, 45, 50, 55, 60, 65]
)

base_columns = [workclass, education, marital_status, relationship, occupation, age_bucket]

crossed_columns = [
    tf.feature_column.crossed_column(
        keys=["education", "occupation"], hash_bucket_size=1000
    ),
    tf.feature_column.crossed_column(
        keys=[age_bucket, "education", "occupation"], hash_bucket_size=1000
    )
]

deep_columns = [
    age,
    education_num,
    capital_gain,
    capital_loss,
    hours_per_week,
    tf.feature_column.indicator_column(workclass),  # 做one-hot，然后送入dnn layer
    tf.feature_column.indicator_column(education),
    tf.feature_column.indicator_column(marital_status),
    tf.feature_column.indicator_column(relationship),

    # 展示embedding的使用
    tf.feature_column.embedding_column(occupation, dimension=8)
]

# 定义数据
_CSV_COLUMNS = [
    "age", "workclass", "fnlwgt", "education", "education_num",
    "marital_status", "occupation", "relationship", "race", "sex",
    "capital_gain", "capital_loss", "hours_per_week", "native_country", "label"
]
_CSV_COLUMN_DEFAULTS = [
    [0], [''], [0], [''], [0],
    [''], [''], [''], [''], [''],
    [0], [0], [0], [''], ['']
]

_NUM_EXAMPLES = {
    "train": 32561,
    "validation": 16281
}

def create_model():
    model = tf.estimator.DNNLinearCombinedClassifier(
        model_dir="./model/wd/",
        linear_feature_columns=base_columns + crossed_columns,
        dnn_feature_columns=deep_columns,
        dnn_hidden_units=[100, 50],
        linear_optimizer="Ftrl",
        dnn_optimizer="Adagrad",
        n_classes=2,
        batch_norm=False
    )
    return model

def input_fn(data_file, num_epochs, shuffle, batch_size):
    """为Estimator创建一个input function"""
    assert tf.io.gfile.GFile(data_file), "{0} not found.".format(data_file)

    def parse_csv(line):
        # tf.decode_csv会把csv文件转换成 a list of Tensor,一列一个
        # record_defaults用于指明每一列的缺失值用什么填充
        columns = tf.io.decode_csv(line, record_defaults=_CSV_COLUMN_DEFAULTS)
        features = dict(zip(_CSV_COLUMNS, columns))
        labels = features.pop('label')
        # tf.equal(x, y) 返回一个bool类型Tensor， 表示x == y, element-wise
        # 注意数据重的空格
        return features, tf.equal(labels, ' >50K')

    dataset = tf.data.TextLineDataset(data_file).map(parse_csv, num_parallel_calls=5)

    if shuffle:
        dataset = dataset.shuffle(buffer_size=_NUM_EXAMPLES['train'] + _NUM_EXAMPLES['validation'])

    dataset = dataset.repeat(num_epochs)
    dataset = dataset.batch(batch_size)

    return dataset

if __name__ == "__main__":
    train_epochs = 20
    batch_size = 256

    model = create_model()
    for n in range(train_epochs):
        print("train model start ...")
        model.train(input_fn=lambda: input_fn(train_data, train_epochs, True, batch_size))
        predict_results = model.predict(input_fn=lambda: input_fn(test_data, train_epochs, False, batch_size))

        print("test model start ...")
        results = model.evaluate(input_fn=lambda: input_fn(test_data, train_epochs, False, batch_size))
        # print(results)

        print('{0:-^30}'.format('evaluate at epoch %d' % ((n + 1))))
        # results 是一个字典
        print(pd.Series(results).to_frame('values'))

-----evaluate at epoch 20-----
                            values
accuracy                  0.826301
accuracy_baseline         0.763774
auc                       0.852878
auc_precision_recall      0.686778
average_loss              0.381445
label/mean                0.236226
loss                      0.381446
precision                 0.727232
prediction/mean           0.249888
recall                    0.423557
global_step           51008.000000

Process finished with exit code 0