Esempio n. 1
0
def main(_):
    flags_obj = tf.flags.FLAGS
    euler_graph = tf_euler.dataset.get_dataset(flags_obj.dataset)
    euler_graph.load_graph()

    fanouts = list(map(int, flags_obj.fanouts))
    assert flags_obj.layers == len(fanouts)
    dims = [flags_obj.hidden_dim] * (flags_obj.layers + 1)
    if flags_obj.run_mode == 'train':
        metapath = [euler_graph.train_edge_type] * flags_obj.layers
    else:
        metapath = [euler_graph.all_edge_type] * flags_obj.layers
    num_steps = int((euler_graph.total_size + 1) // flags_obj.batch_size *
                    flags_obj.num_epochs)

    model = get_solution_model(conv=flags_obj.conv,
                               dataflow=flags_obj.flow,
                               dims=dims,
                               fanouts=fanouts,
                               metapath=metapath,
                               feature_idx=euler_graph.feature_idx,
                               feature_dim=euler_graph.feature_dim,
                               label_idx=euler_graph.label_idx,
                               label_dim=euler_graph.label_dim)

    params = {
        'train_node_type': euler_graph.train_node_type[0],
        'batch_size': flags_obj.batch_size,
        'optimizer': flags_obj.optimizer,
        'learning_rate': flags_obj.learning_rate,
        'log_steps': flags_obj.log_steps,
        'model_dir': flags_obj.model_dir,
        'id_file': euler_graph.id_file,
        'infer_dir': flags_obj.model_dir,
        'total_step': num_steps
    }
    config = tf.estimator.RunConfig(log_step_count_steps=None)
    model_estimator = NodeEstimator(model, params, config)

    if flags_obj.run_mode == 'train':
        model_estimator.train()
    elif flags_obj.run_mode == 'evaluate':
        model_estimator.evaluate()
    elif flags_obj.run_mode == 'infer':
        model_estimator.infer()
    else:
        raise ValueError('Run mode not exist!')
Esempio n. 2
0
def main(_):
    flags_obj = tf.flags.FLAGS
    euler_graph = tf_euler.dataset.get_dataset(flags_obj.dataset)
    euler_graph.load_graph()

    if flags_obj.run_mode == 'train':
        metapath = [euler_graph.train_edge_type] * flags_obj.layers
    else:
        metapath = [euler_graph.all_edge_type] * flags_obj.layers
    head_num = [flags_obj.num_heads] * flags_obj.layers
    num_steps = int((euler_graph.total_size + 1) // flags_obj.batch_size *
                    flags_obj.num_epochs)

    model = GeniePath(flags_obj.hidden_dim,
                      metapath,
                      euler_graph.label_idx,
                      euler_graph.label_dim,
                      max_id=euler_graph.max_node_id,
                      feature_idx=euler_graph.feature_idx,
                      feature_dim=euler_graph.feature_dim,
                      use_id=flags_obj.use_id,
                      embedding_dim=flags_obj.embedding_dim,
                      head_num=head_num)

    params = {
        'train_node_type': euler_graph.train_node_type[0],
        'batch_size': flags_obj.batch_size,
        'learning_rate': flags_obj.learning_rate,
        'log_steps': flags_obj.log_steps,
        'model_dir': flags_obj.model_dir,
        'id_file': euler_graph.id_file,
        'infer_dir': flags_obj.model_dir,
        'total_step': num_steps
    }
    config = tf.estimator.RunConfig(log_step_count_steps=None)
    model_estimator = NodeEstimator(model, params, config)

    if flags_obj.run_mode == 'train':
        model_estimator.train()
    elif flags_obj.run_mode == 'evaluate':
        model_estimator.evaluate()
    elif flags_obj.run_mode == 'infer':
        model_estimator.infer()
    else:
        raise ValueError('Run mode not exist!')
Esempio n. 3
0
def main(_):
    flags_obj = tf.flags.FLAGS
    euler_graph = tf_euler.dataset.get_dataset(flags_obj.dataset)
    euler_graph.load_graph()

    dims = [flags_obj.hidden_dim, flags_obj.hidden_dim]
    if flags_obj.run_mode == 'train':
        metapath = [euler_graph.train_edge_type]
    else:
        metapath = [euler_graph.all_edge_type]
    num_steps = int((euler_graph.total_size + 1) // flags_obj.batch_size *
                    flags_obj.num_epochs)

    model = ARMA(dims,
                 metapath,
                 euler_graph.feature_idx,
                 euler_graph.feature_dim,
                 euler_graph.label_idx,
                 euler_graph.label_dim,
                 num_layers=flags_obj.layers,
                 K=flags_obj.K)
    params = {
        'train_node_type': euler_graph.train_node_type[0],
        'batch_size': flags_obj.batch_size,
        'optimizer': flags_obj.optimizer,
        'learning_rate': flags_obj.learning_rate,
        'log_steps': flags_obj.log_steps,
        'model_dir': flags_obj.model_dir,
        'id_file': euler_graph.id_file,
        'infer_dir': flags_obj.model_dir,
        'total_size': euler_graph.total_size,
        'total_step': num_steps
    }
    config = tf.estimator.RunConfig(log_step_count_steps=None)
    model_estimator = NodeEstimator(model, params, config)

    if flags_obj.run_mode == 'train':
        model_estimator.train()
    elif flags_obj.run_mode == 'evaluate':
        model_estimator.evaluate()
    elif flags_obj.run_mode == 'infer':
        model_estimator.infer()
    else:
        raise ValueError('Run mode not exist!')
Esempio n. 4
0
def main(_):
    flags_obj = tf.flags.FLAGS
    euler_graph = tf_euler.dataset.get_dataset(flags_obj.dataset)
    euler_graph.load_graph()

    dims = [flags_obj.hidden_dim] * (flags_obj.layers + 1)
    if flags_obj.run_mode == 'train':
        metapath = [euler_graph.all_edge_type] * flags_obj.layers
    else:
        metapath = [euler_graph.all_edge_type] * flags_obj.layers
    num_steps = int((euler_graph.total_size + 1) // flags_obj.batch_size *
                    flags_obj.num_epochs)

    model = UnsupervisedRGCN(euler_graph.all_node_type,
                             euler_graph.all_edge_type,
                             euler_graph.max_node_id, dims, metapath,
                             euler_graph.max_edge_id, euler_graph.edge_id_idx,
                             euler_graph.edge_id_dim, flags_obj.embedding_dim,
                             flags_obj.num_negs, flags_obj.metric)

    params = {
        'train_node_type': euler_graph.train_node_type[0],
        'batch_size': flags_obj.batch_size,
        'optimizer': flags_obj.optimizer,
        'learning_rate': flags_obj.learning_rate,
        'log_steps': flags_obj.log_steps,
        'model_dir': flags_obj.model_dir,
        'id_file': euler_graph.node_id_file,
        'infer_dir': flags_obj.model_dir,
        'total_step': num_steps
    }
    config = tf.estimator.RunConfig(log_step_count_steps=None)
    model_estimator = NodeEstimator(model, params, config)

    if flags_obj.run_mode == 'train':
        model_estimator.train()
    elif flags_obj.run_mode == 'evaluate':
        model_estimator.evaluate()
    elif flags_obj.run_mode == 'infer':
        model_estimator.infer()
    else:
        raise ValueError('Run mode not exist!')
Esempio n. 5
0
def main(_):
    flags_obj = tf.flags.FLAGS
    euler_graph = tf_euler.dataset.get_dataset(flags_obj.dataset)
    euler_graph.load_graph()
    num_steps = int((euler_graph.total_size + 1) // flags_obj.batch_size *
                    flags_obj.num_epochs)

    model = Line(euler_graph.all_node_type,
                 euler_graph.all_edge_type,
                 euler_graph.max_node_id,
                 flags_obj.hidden_dim,
                 num_negs=flags_obj.num_negs,
                 order=flags_obj.order,
                 use_id=flags_obj.use_id,
                 metric=flags_obj.metric,
                 embedding_dim=flags_obj.embedding_dim)

    params = {
        'train_node_type': euler_graph.all_node_type,
        'batch_size': flags_obj.batch_size,
        'optimizer': flags_obj.optimizer,
        'learning_rate': flags_obj.learning_rate,
        'log_steps': flags_obj.log_steps,
        'model_dir': flags_obj.model_dir,
        'id_file': euler_graph.id_file,
        'infer_dir': flags_obj.model_dir,
        'total_step': num_steps
    }
    config = tf.estimator.RunConfig(log_step_count_steps=None)
    model_estimator = NodeEstimator(model, params, config)

    if flags_obj.run_mode == 'train':
        model_estimator.train()
    elif flags_obj.run_mode == 'evaluate':
        model_estimator.evaluate()
    elif flags_obj.run_mode == 'infer':
        model_estimator.infer()
    else:
        raise ValueError('Run mode not exist!')