Hetu/examples/ctr/run_tf_local.py

import numpy as np
import tensorflow as tf
import time
import argparse
from tqdm import tqdm
from sklearn import metrics


def train_criteo(model, args):
    if args.all:
        from models.load_data import process_all_criteo_data
        dense, sparse, all_labels = process_all_criteo_data()
        dense_feature, val_dense = dense
        sparse_feature, val_sparse = sparse
        labels, val_labels = all_labels
    else:
        from models.load_data import process_sampled_criteo_data
        dense_feature, sparse_feature, labels = process_sampled_criteo_data()

    batch_size = 128
    dense_input = tf.compat.v1.placeholder(tf.float32, [batch_size, 13])
    sparse_input = tf.compat.v1.placeholder(tf.int32, [batch_size, 26])
    y_ = y_ = tf.compat.v1.placeholder(tf.float32, [batch_size, 1])

    loss, y, opt = model(dense_input, sparse_input, y_)
    train_op = opt.minimize(loss)

    init = tf.compat.v1.global_variables_initializer()
    gpu_options = tf.compat.v1.GPUOptions(allow_growth=True)
    sess = tf.compat.v1.Session(
        config=tf.compat.v1.ConfigProto(gpu_options=gpu_options))
    sess.run(init)

    my_feed_dict = {
        dense_input: np.empty(shape=(batch_size, 13)),
        sparse_input: np.empty(shape=(batch_size, 26)),
        y_: np.empty(shape=(batch_size, 1)),
    }

    if args.all:
        raw_log_file = './logs/tf_local_%s.log' % (args.model)
        print('Processing all data, log to', raw_log_file)
        log_file = open(raw_log_file, 'w')
        iterations = dense_feature.shape[0] // batch_size
        total_epoch = 11
        start_index = 0
        for ep in range(total_epoch):
            print("epoch %d" % ep)
            st_time = time.time()
            train_loss, train_acc, train_auc = [], [], []
            for it in tqdm(range(iterations // 10 + (ep % 10 == 9) * (iterations % 10))):
                my_feed_dict[dense_input][:] = dense_feature[start_index: start_index + batch_size]
                my_feed_dict[sparse_input][:] = sparse_feature[start_index: start_index + batch_size]
                my_feed_dict[y_][:] = labels[start_index: start_index+batch_size]
                start_index += batch_size
                if start_index + batch_size > dense_feature.shape[0]:
                    start_index = 0
                loss_val = sess.run([loss, y, y_, train_op],
                                    feed_dict=my_feed_dict)
                pred_val = loss_val[1]
                true_val = loss_val[2]
                acc_val = np.equal(
                    true_val,
                    pred_val > 0.5)
                train_loss.append(loss_val[0])
                train_acc.append(acc_val)
                train_auc.append(metrics.roc_auc_score(true_val, pred_val))
            tra_accuracy = np.mean(train_acc)
            tra_loss = np.mean(train_loss)
            tra_auc = np.mean(train_auc)
            en_time = time.time()
            train_time = en_time - st_time
            printstr = "train_loss: %.4f, train_acc: %.4f, train_auc: %.4f, train_time: %.4f"\
                % (tra_loss, tra_accuracy, tra_auc, train_time)
            print(printstr)
            log_file.write(printstr + '\n')
            log_file.flush()

    else:
        iteration = dense_feature.shape[0] // batch_size

        epoch = 50
        for ep in range(epoch):
            print('epoch', ep)
            if ep == 5:
                start = time.time()
            ep_st = time.time()
            train_loss = []
            train_acc = []
            for idx in range(iteration):
                start_index = idx * batch_size
                my_feed_dict[dense_input][:] = dense_feature[start_index: start_index + batch_size]
                my_feed_dict[sparse_input][:] = sparse_feature[start_index: start_index + batch_size]
                my_feed_dict[y_][:] = labels[start_index: start_index+batch_size]

                loss_val = sess.run([loss, y, y_, train_op],
                                    feed_dict=my_feed_dict)
                pred_val = loss_val[1]
                true_val = loss_val[2]
                if pred_val.shape[1] == 1:  # for criteo case
                    acc_val = np.equal(
                        true_val,
                        pred_val > 0.5)
                else:
                    acc_val = np.equal(
                        np.argmax(pred_val, 1),
                        np.argmax(true_val, 1)).astype(np.float32)
                train_loss.append(loss_val[0])
                train_acc.append(acc_val)
            tra_accuracy = np.mean(train_acc)
            tra_loss = np.mean(train_loss)
            ep_en = time.time()
            print("train_loss: %.4f, train_acc: %.4f, train_time: %.4f"
                  % (tra_loss, tra_accuracy, ep_en - ep_st))
        print('all time:', (time.time() - start))


def train_adult(model):
    batch_size = 128
    total_epoch = 50
    dim_wide = 809

    X_deep = []
    for i in range(8):
        X_deep.append(tf.compat.v1.placeholder(tf.int32, [batch_size, 1]))
    for i in range(4):
        X_deep.append(tf.compat.v1.placeholder(tf.float32, [batch_size, 1]))
    X_wide = tf.compat.v1.placeholder(tf.float32, [batch_size, dim_wide])
    y_ = tf.compat.v1.placeholder(tf.float32, [batch_size, 2])
    loss, y, train_op = model(X_deep, X_wide, y_)

    init = tf.global_variables_initializer()

    gpu_options = tf.GPUOptions(allow_growth=True)
    sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

    sess.run(init)

    from models.load_data import load_adult_data
    x_train_deep, x_train_wide, y_train = load_adult_data(return_val=False)

    iterations = x_train_deep.shape[0] // batch_size
    for ep in range(total_epoch):
        print('epoch', ep)
        if ep == 5:
            start = time.time()
        ep_st = time.time()
        train_loss = []
        train_acc = []
        pre_index = 0

        for it in range(iterations):
            batch_x_deep = x_train_deep[pre_index:pre_index + batch_size]
            batch_x_wide = x_train_wide[pre_index:pre_index + batch_size]
            batch_y = y_train[pre_index:pre_index + batch_size]
            pre_index += batch_size

            my_feed_dict = dict()
            for i in range(12):
                my_feed_dict[X_deep[i]] = np.array(
                    batch_x_deep[:, 1]).reshape(-1, 1)

            my_feed_dict[X_wide] = np.array(batch_x_wide)
            my_feed_dict[y_] = batch_y
            loss_val = sess.run([loss, y, y_, train_op],
                                feed_dict=my_feed_dict)
            acc_val = np.equal(
                np.argmax(loss_val[1], 1),
                np.argmax(loss_val[2], 1)).astype(np.float32)
            train_loss.append(loss_val[0])
            train_acc.append(acc_val)
        tra_accuracy = np.mean(train_acc)
        tra_loss = np.mean(train_loss)
        ep_en = time.time()
        print("train_loss: %.4f, train_acc: %.4f, train_time: %.4f"
              % (tra_loss, tra_accuracy, ep_en - ep_st))
    print('all time:', (time.time() - start))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, required=True,
                        help="model to be tested")
    parser.add_argument("--all", action="store_true",
                        help="whether to use all data")
    args = parser.parse_args()
    raw_model = args.model
    import tf_models
    model = eval('tf_models.' + raw_model)
    dataset = raw_model.split('_')[-1]
    print('Model:', raw_model)

    if dataset == 'criteo':
        train_criteo(model, args)
    elif dataset == 'adult':
        train_adult(model)
    else:
        raise NotImplementedError


if __name__ == '__main__':
    main()