convert_tfrecord.py 文件源码

def main(unused_argv):

    # Get the data.
    data_train = np.loadtxt(os.path.join(FLAGS.input_directory,'train'), delimiter=',')
    data_test = np.loadtxt(os.path.join(FLAGS.input_directory, 'test'), delimiter=',')

    X_train, X_val, y_train, y_val = train_test_split(data_train[:,1:], data_train[:,0].astype(np.int32),
                                                          test_size=FLAGS.validation_ratio,
                                                          random_state=100)
    X_test = data_test[:, 1:]
    y_test = data_test[:, 0].astype(np.int32)

    # Convert to Examples and write the result to TFRecords.
    convert_to((X_train, y_train), PREFIX + '_train')
    convert_to((X_val, y_val), PREFIX + '_validation')
    convert_to((X_test, y_test), PREFIX + '_test')