GFNCP: Consistent Amortized Clustering via Generative Flow Networks [AISTATS 2025]

Irit Chelly, Roy Uziel, Oren Freifeld and Ari Pakman.

Pytorch implementation of GFNCP.

Requirements

python 3
torch
numpy
wandb

How to use

To run the code:

python main.py --dataset MNIST --data_path '/your_path'

To use a pretrained model, use the load-model flag, and the latest checkpoint from saved_models folder will be used.

For tracking loss and metrics values during training and evaluation, use wandb flag to log these values to Weights and Biases.

Make sure to update wandb entity and project names in main.py:

def init_wandb(args, params):
    if has_wandb:
        wnb = wandb.init(entity='your_entity', project='your_project', name='experiment_name', config=args, settings=wandb.Settings(_service_wait=300))
        ...

Use params.py for setting batch_size, iterations number, and other hyper-parameters.

License

This project is released under the MIT license. Please see the LICENSE file for more information.

Citation

If you find this repository helpful, please consider citing our paper:

TBD