| 29 | |
| 30 | |
| 31 | def define_network_flags(): |
| 32 | tf.flags.DEFINE_string('dataset', 'cora', 'Dataset name.') |
| 33 | tf.flags.DEFINE_integer('hidden_dim', 32, 'Hidden dimension.') |
| 34 | tf.flags.DEFINE_integer('layers', 2, 'SAGE convolution layer number.') |
| 35 | tf.flags.DEFINE_enum('conv', 'sage', ['sage', 'gcn'], 'GNN Convolution.') |
| 36 | tf.flags.DEFINE_enum('flow', 'sage', ['sage', 'full'], 'GNN dataflow.') |
| 37 | tf.flags.DEFINE_list('fanouts', [10, 10], 'GraphSage fanouts.') |
| 38 | tf.flags.DEFINE_integer('batch_size', 32, 'Mini-batch size') |
| 39 | tf.flags.DEFINE_integer('num_epochs', 10, 'Epochs to train') |
| 40 | tf.flags.DEFINE_integer('log_steps', 20, 'Number of steps to print log.') |
| 41 | tf.flags.DEFINE_string('model_dir', 'ckpt', 'Model checkpoint.') |
| 42 | tf.flags.DEFINE_float('learning_rate', 0.01, 'Learning rate.') |
| 43 | tf.flags.DEFINE_enum('optimizer', 'adam', ['adam', 'adagrad', 'sgd', 'momentum'], |
| 44 | 'Optimizer algorithm') |
| 45 | tf.flags.DEFINE_enum('run_mode', 'train', ['train', 'evaluate', 'infer'], |
| 46 | 'Run mode.') |
| 47 | |
| 48 | |
| 49 | def main(_): |