PopGen is a generative modelling toolkit written in PyTorch. It aims to provide high quality reference implementations and reusable components, with an emphasis on likelihood based models and representation learning.
Experiment Framework
Experiments in PopGen are organised around three core concepts.
Models
Any PyTorch model (inherits from nn.Module)
Should not contain optimisation related code.
Must be included models/__init__.py to be dynamically loaded by the experiment runner.
Refer to models/vae.py for an example.
Workers
Responsible for training and evaluating models
Must implement .train and .evaluate
Should inherit from the AbstractWorker class, which will:
provide utility methods for saving and loading experiment checkpoints.
enforce a common training interface
Must be included in workers/__init__.py to be dynamically loaded by the experiment runner.
Refer to workers/vae_worker.py for an example.
Datasets
Must implement __len__ and __getitem__
Will be wrapped with the PyTorch DataLoader class for batching and collation
Must be included in datasets/__init__.py to be dynamically loaded by the experiment runner.
Refer to datasets/binary_mnist.py for an example.
Note that there is not a strict 1:1 relationship between workers and models, models and datasets etc. It is up to the developer to ensure a compatible API between these various components.