Exemple #1
0
def test_save_load():
    data = pd.DataFrame({
        'continuous': np.random.random(100),
        'discrete': np.random.choice(['a', 'b', 'c'], 100)
    })
    discrete_columns = ['discrete']

    ctgan = CTGANSynthesizer(epochs=1)
    ctgan.fit(data, discrete_columns)

    with tf.TemporaryDirectory() as temporary_directory:
        ctgan.save(temporary_directory + "test_tvae.pkl")
        ctgan = CTGANSynthesizer.load(temporary_directory + "test_tvae.pkl")

    sampled = ctgan.sample(1000)
    assert set(sampled.columns) == {'continuous', 'discrete'}
    assert set(sampled['discrete'].unique()) == {'a', 'b', 'c'}
Exemple #2
0
def main():
    """CLI."""
    args = _parse_args()
    if args.tsv:
        data, discrete_columns = read_tsv(args.data, args.metadata)
    else:
        data, discrete_columns = read_csv(args.data, args.metadata,
                                          args.header, args.discrete)

    if args.load:
        model = CTGANSynthesizer.load(args.load)
    else:
        generator_dim = [int(x) for x in args.generator_dim.split(',')]
        discriminator_dim = [int(x) for x in args.discriminator_dim.split(',')]
        model = CTGANSynthesizer(embedding_dim=args.embedding_dim,
                                 generator_dim=generator_dim,
                                 discriminator_dim=discriminator_dim,
                                 generator_lr=args.generator_lr,
                                 generator_decay=args.generator_decay,
                                 discriminator_lr=args.discriminator_lr,
                                 discriminator_decay=args.discriminator_decay,
                                 batch_size=args.batch_size,
                                 epochs=args.epochs)
    model.fit(data, discrete_columns)

    if args.save is not None:
        model.save(args.save)

    num_samples = args.num_samples or len(data)

    if args.sample_condition_column is not None:
        assert args.sample_condition_column_value is not None

    sampled = model.sample(num_samples, args.sample_condition_column,
                           args.sample_condition_column_value)

    if args.tsv:
        write_tsv(sampled, args.metadata, args.output)
    else:
        sampled.to_csv(args.output, index=False)