Add files using upload-large-folder tool

conditional-flow-matching/.github/workflows/code-quality.yaml

@@ -0,0 +1,26 @@
+# Same as `code-quality-pr.yaml` but triggered on commit to main branch
+# and runs on all files (instead of only the changed ones)
+
+name: Code Quality Main
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main, "release/*"]
+
+jobs:
+  code-quality:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v5
+
+      - name: Set up Python
+        uses: actions/setup-python@v6
+        with:
+          python-version: "3.13"
+
+      - name: Run pre-commits
+        uses: pre-commit/action@v3.0.1

conditional-flow-matching/.github/workflows/python-publish.yml

@@ -0,0 +1,38 @@
+# This workflow will upload a Python Package using Twine when a release is created
+# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python#publishing-to-package-registries
+
+# This workflow uses actions that are not certified by GitHub.
+# They are provided by a third-party and are governed by
+# separate terms of service, privacy policy, and support
+# documentation.
+
+name: Upload Python Package
+
+on:
+  release:
+    types: [published]
+
+permissions:
+  contents: read
+
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python
+        uses: actions/setup-python@v3
+        with:
+          python-version: "3.x"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install build
+      - name: Build package
+        run: python -m build
+      - name: Publish package
+        uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
+        with:
+          user: __token__
+          password: ${{ secrets.PYPI_API_TOKEN }}

conditional-flow-matching/.github/workflows/test.yaml

@@ -0,0 +1,71 @@
+name: TorchCFM Tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main, "release/*"]
+
+jobs:
+  run_tests_ubuntu:
+    runs-on: ${{ matrix.os }}
+
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, macos-latest, windows-latest]
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v5
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v6
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install pytest
+          pip install sh
+          pip install -e .
+
+      - name: List dependencies
+        run: |
+          python -m pip list
+
+      - name: Run pytest
+        run: |
+          pytest -v --ignore=examples --ignore=runner
+
+  # upload code coverage report
+  code-coverage-torchcfm:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v5
+
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v6
+        with:
+          python-version: "3.10"
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install pytest
+          pip install pytest-cov[toml]
+          pip install sh
+          pip install -e .
+
+      - name: Run tests and collect coverage
+        run: pytest . --cov torchcfm --ignore=runner --ignore=examples --ignore=torchcfm/models/ --cov-fail-under=30
+
+      - name: Upload coverage to Codecov
+        uses: codecov/codecov-action@v3
+        with:
+          name: codecov-torchcfm
+          verbose: true

conditional-flow-matching/.github/workflows/test_runner.yaml

@@ -0,0 +1,75 @@
+name: Runner Tests
+
+#on:
+#  push:
+#    branches: [main]
+#  pull_request:
+#    branches: [main, "release/*"]
+
+jobs:
+  run_tests_ubuntu:
+    runs-on: ${{ matrix.os }}
+
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, ubuntu-20.04, macos-latest, windows-latest]
+        python-version: ["3.9", "3.10"]
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v5
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v6
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: |
+          # Fix pip version < 24.1 due to lightning incomaptibility
+          python -m pip install pip==23.2.1
+          pip install -r runner-requirements.txt
+          pip install pytest
+          pip install sh
+          pip install -e .
+
+      - name: List dependencies
+        run: |
+          python -m pip list
+
+      - name: Run pytest
+        run: |
+          pytest -v runner
+
+  # upload code coverage report
+  code-coverage-runner:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v5
+
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v6
+        with:
+          python-version: "3.10"
+
+      - name: Install dependencies
+        run: |
+          # Fix pip version < 24.1 due to lightning incomaptibility
+          python -m pip install pip==23.2.1
+          pip install -r runner-requirements.txt
+          pip install pytest
+          pip install pytest-cov[toml]
+          pip install sh
+          pip install -e .
+
+      - name: Run tests and collect coverage
+        run: pytest runner --cov runner --cov-fail-under=30 # NEEDS TO BE UPDATED WHEN CHANGING THE NAME OF "src" FOLDER
+
+      - name: Upload coverage to Codecov
+        uses: codecov/codecov-action@v3
+        with:
+          name: codecov-runner
+          verbose: true

conditional-flow-matching/runner/configs/callbacks/default.yaml

@@ -0,0 +1,22 @@
+defaults:
+  - model_checkpoint.yaml
+  - early_stopping.yaml
+  - model_summary.yaml
+  - rich_progress_bar.yaml
+  - _self_
+
+model_checkpoint:
+  dirpath: ${paths.output_dir}/checkpoints
+  filename: "epoch_{epoch:04d}"
+  monitor: "val/loss"
+  mode: "min"
+  save_last: True
+  auto_insert_metric_name: False
+
+early_stopping:
+  monitor: "val/loss"
+  patience: 100
+  mode: "min"
+
+model_summary:
+  max_depth: -1

conditional-flow-matching/runner/configs/callbacks/early_stopping.yaml

@@ -0,0 +1,17 @@
+# https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.callbacks.EarlyStopping.html
+
+# Monitor a metric and stop training when it stops improving.
+# Look at the above link for more detailed information.
+early_stopping:
+  _target_: pytorch_lightning.callbacks.EarlyStopping
+  monitor: ??? # quantity to be monitored, must be specified !!!
+  min_delta: 0. # minimum change in the monitored quantity to qualify as an improvement
+  patience: 3 # number of checks with no improvement after which training will be stopped
+  verbose: False # verbosity mode
+  mode: "min" # "max" means higher metric value is better, can be also "min"
+  strict: True # whether to crash the training if monitor is not found in the validation metrics
+  check_finite: True # when set True, stops training when the monitor becomes NaN or infinite
+  stopping_threshold: null # stop training immediately once the monitored quantity reaches this threshold
+  divergence_threshold: null # stop training as soon as the monitored quantity becomes worse than this threshold
+  check_on_train_epoch_end: null # whether to run early stopping at the end of the training epoch
+  # log_rank_zero_only: False  # this keyword argument isn't available in stable version

conditional-flow-matching/runner/configs/callbacks/model_checkpoint.yaml

@@ -0,0 +1,19 @@
+# https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.callbacks.ModelCheckpoint.html
+
+# Save the model periodically by monitoring a quantity.
+# Look at the above link for more detailed information.
+model_checkpoint:
+  _target_: pytorch_lightning.callbacks.ModelCheckpoint
+  dirpath: null # directory to save the model file
+  filename: null # checkpoint filename
+  monitor: null # name of the logged metric which determines when model is improving
+  verbose: False # verbosity mode
+  save_last: null # additionally always save an exact copy of the last checkpoint to a file last.ckpt
+  save_top_k: 1 # save k best models (determined by above metric)
+  mode: "min" # "max" means higher metric value is better, can be also "min"
+  auto_insert_metric_name: True # when True, the checkpoints filenames will contain the metric name
+  save_weights_only: False # if True, then only the model’s weights will be saved
+  every_n_train_steps: null # number of training steps between checkpoints
+  train_time_interval: null # checkpoints are monitored at the specified time interval
+  every_n_epochs: null # number of epochs between checkpoints
+  save_on_train_epoch_end: null # whether to run checkpointing at the end of the training epoch or the end of validation

conditional-flow-matching/runner/configs/callbacks/model_summary.yaml

@@ -0,0 +1,7 @@
+# https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.callbacks.RichModelSummary.html
+
+# Generates a summary of all layers in a LightningModule with rich text formatting.
+# Look at the above link for more detailed information.
+model_summary:
+  _target_: pytorch_lightning.callbacks.RichModelSummary
+  max_depth: 1 # the maximum depth of layer nesting that the summary will include

conditional-flow-matching/runner/configs/callbacks/no_stopping.yaml

@@ -0,0 +1,15 @@
+defaults:
+  - model_checkpoint.yaml
+  - model_summary.yaml
+  - rich_progress_bar.yaml
+  - _self_
+
+model_checkpoint:
+  dirpath: ${paths.output_dir}/checkpoints
+  filename: "epoch_{epoch:04d}"
+  save_last: True
+  every_n_epochs: 100 # number of epochs between checkpoints
+  auto_insert_metric_name: False
+
+model_summary:
+  max_depth: 3

conditional-flow-matching/runner/configs/callbacks/rich_progress_bar.yaml

@@ -0,0 +1,6 @@
+# https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.callbacks.RichProgressBar.html
+
+# Create a progress bar with rich text formatting.
+# Look at the above link for more detailed information.
+rich_progress_bar:
+  _target_: pytorch_lightning.callbacks.RichProgressBar

conditional-flow-matching/runner/configs/datamodule/cifar.yaml

@@ -0,0 +1,11 @@
+_target_: src.datamodules.cifar10_datamodule.CIFAR10DataModule
+#_target_: pl_bolts.datamodules.cifar10_datamodule.CIFAR10DataModule
+data_dir: ${paths.data_dir}
+batch_size: 128
+val_split: 0.0
+num_workers: 0
+normalize: True
+seed: 42
+shuffle: True
+pin_memory: True
+drop_last: False

conditional-flow-matching/runner/configs/datamodule/custom_dist.yaml

@@ -0,0 +1,17 @@
+_target_: src.datamodules.distribution_datamodule.TrajectoryNetDistributionTrajectoryDataModule
+#_target_: src.datamodules.distribution_datamodule.DistributionDataModule
+
+data_dir: ${paths.data_dir} # data_dir is specified in config.yaml
+train_val_test_split: 1000
+batch_size: 100
+num_workers: 0
+pin_memory: False
+
+system: ${paths.data_dir}/embryoid_anndata_small_v2.h5ad
+
+system_kwargs:
+  max_dim: 1e10
+  embedding_name: "phate"
+  #embedding_name: "highly_variable"
+  whiten: True
+  #whiten: False

conditional-flow-matching/runner/configs/datamodule/eb_full.yaml

@@ -0,0 +1,14 @@
+# https://github.com/DiffEqML/torchdyn/blob/master/torchdyn/datasets/static_datasets.py
+_target_: src.datamodules.distribution_datamodule.TorchDynDataModule
+#_target_: src.datamodules.distribution_datamodule.DistributionDataModule
+
+train_val_test_split: [0.8, 0.1, 0.1]
+batch_size: 128
+num_workers: 0
+pin_memory: False
+
+system: ${paths.data_dir}/eb_velocity_v5.npz
+
+system_kwargs:
+  max_dim: 100
+  whiten: False

conditional-flow-matching/runner/configs/datamodule/funnel.yaml

@@ -0,0 +1,12 @@
+_target_: src.datamodules.distribution_datamodule.TorchDynDataModule
+#_target_: src.datamodules.distribution_datamodule.DistributionDataModule
+
+train_val_test_split: [10000, 1000, 1000]
+batch_size: 128
+num_workers: 0
+pin_memory: False
+
+system: "funnel"
+
+system_kwargs:
+  dim: 10

conditional-flow-matching/runner/configs/datamodule/gaussians.yaml

@@ -0,0 +1,13 @@
+# https://github.com/DiffEqML/torchdyn/blob/master/torchdyn/datasets/static_datasets.py
+_target_: src.datamodules.distribution_datamodule.TorchDynDataModule
+#_target_: src.datamodules.distribution_datamodule.DistributionDataModule
+
+train_val_test_split: [10000, 1000, 1000]
+batch_size: 128
+num_workers: 0
+pin_memory: False
+
+system: "gaussians"
+
+system_kwargs:
+  noise: 1e-4

conditional-flow-matching/runner/configs/datamodule/moons.yaml

@@ -0,0 +1,10 @@
+# https://github.com/DiffEqML/torchdyn/blob/master/torchdyn/datasets/static_datasets.py
+_target_: src.datamodules.distribution_datamodule.SKLearnDataModule
+#_target_: src.datamodules.distribution_datamodule.DistributionDataModule
+
+train_val_test_split: [10000, 1000, 1000]
+batch_size: 128
+num_workers: 0
+pin_memory: False
+
+system: "moons"

conditional-flow-matching/runner/configs/datamodule/scurve.yaml

@@ -0,0 +1,10 @@
+# https://github.com/DiffEqML/torchdyn/blob/master/torchdyn/datasets/static_datasets.py
+_target_: src.datamodules.distribution_datamodule.SKLearnDataModule
+#_target_: src.datamodules.distribution_datamodule.DistributionDataModule
+
+train_val_test_split: [10000, 1000, 1000]
+batch_size: 128
+num_workers: 0
+pin_memory: False
+
+system: "scurve"

conditional-flow-matching/runner/configs/datamodule/sklearn.yaml

@@ -0,0 +1,10 @@
+# https://github.com/DiffEqML/torchdyn/blob/master/torchdyn/datasets/static_datasets.py
+_target_: src.datamodules.distribution_datamodule.SKLearnDataModule
+#_target_: src.datamodules.distribution_datamodule.DistributionDataModule
+
+train_val_test_split: [10000, 1000, 1000]
+batch_size: 128
+num_workers: 0
+pin_memory: False
+
+system: "scurve"

conditional-flow-matching/runner/configs/datamodule/time_dist.yaml

@@ -0,0 +1,12 @@
+_target_: src.datamodules.distribution_datamodule.CustomTrajectoryDataModule
+#_target_: src.datamodules.distribution_datamodule.DistributionDataModule
+
+data_dir: ${paths.data_dir} # data_dir is specified in config.yaml
+train_val_test_split: [0.8, 0.1, 0.1]
+batch_size: 128
+num_workers: 0
+pin_memory: False
+max_dim: 5
+whiten: True
+
+system: ${paths.data_dir}/eb_velocity_v5.npz

conditional-flow-matching/runner/configs/datamodule/torchdyn.yaml

@@ -0,0 +1,13 @@
+# https://github.com/DiffEqML/torchdyn/blob/master/torchdyn/datasets/static_datasets.py
+_target_: src.datamodules.distribution_datamodule.TorchDynDataModule
+#_target_: src.datamodules.distribution_datamodule.DistributionDataModule
+
+train_val_test_split: [10000, 1000, 1000]
+batch_size: 128
+num_workers: 0
+pin_memory: False
+
+system: "moons"
+
+system_kwargs:
+  noise: 1e-4

conditional-flow-matching/runner/configs/datamodule/tree.yaml

@@ -0,0 +1,8 @@
+_target_: src.datamodules.distribution_datamodule.DistributionDataModule
+
+data_dir: ${data_dir} # data_dir is specified in config.yaml
+train_val_test_split: 1000
+batch_size: 100
+num_workers: 0
+pin_memory: False
+p: 2

conditional-flow-matching/runner/configs/datamodule/twodim.yaml

@@ -0,0 +1,8 @@
+_target_: src.datamodules.distribution_datamodule.TwoDimDataModule
+
+train_val_test_split: [10000, 1000, 1000]
+batch_size: 128
+num_workers: 0
+pin_memory: False
+
+system: "moon-8gaussians"

conditional-flow-matching/runner/configs/debug/default.yaml

@@ -0,0 +1,35 @@
+# @package _global_
+
+# default debugging setup, runs 1 full epoch
+# other debugging configs can inherit from this one
+
+# overwrite task name so debugging logs are stored in separate folder
+task_name: "debug"
+
+# disable callbacks and loggers during debugging
+callbacks: null
+logger: null
+
+extras:
+  ignore_warnings: False
+  enforce_tags: False
+
+# sets level of all command line loggers to 'DEBUG'
+# https://hydra.cc/docs/tutorials/basic/running_your_app/logging/
+hydra:
+  job_logging:
+    root:
+      level: DEBUG
+
+  # use this to also set hydra loggers to 'DEBUG'
+  # verbose: True
+
+trainer:
+  max_epochs: 1
+  accelerator: cpu # debuggers don't like gpus
+  devices: 1 # debuggers don't like multiprocessing
+  detect_anomaly: true # raise exception if NaN or +/-inf is detected in any tensor
+
+datamodule:
+  num_workers: 0 # debuggers don't like multiprocessing
+  pin_memory: False # disable gpu memory pin

conditional-flow-matching/runner/configs/debug/fdr.yaml

@@ -0,0 +1,9 @@
+# @package _global_
+
+# runs 1 train, 1 validation and 1 test step
+
+defaults:
+  - default.yaml
+
+trainer:
+  fast_dev_run: true

conditional-flow-matching/runner/configs/debug/limit.yaml

@@ -0,0 +1,12 @@
+# @package _global_
+
+# uses only 1% of the training data and 5% of validation/test data
+
+defaults:
+  - default.yaml
+
+trainer:
+  max_epochs: 3
+  limit_train_batches: 0.01
+  limit_val_batches: 0.05
+  limit_test_batches: 0.05

conditional-flow-matching/runner/configs/debug/overfit.yaml

@@ -0,0 +1,13 @@
+# @package _global_
+
+# overfits to 3 batches
+
+defaults:
+  - default.yaml
+
+trainer:
+  max_epochs: 20
+  overfit_batches: 3
+
+# model ckpt and early stopping need to be disabled during overfitting
+callbacks: null

conditional-flow-matching/runner/configs/debug/profiler.yaml

@@ -0,0 +1,12 @@
+# @package _global_
+
+# runs with execution time profiling
+
+defaults:
+  - default.yaml
+
+trainer:
+  max_epochs: 1
+  profiler: "simple"
+  # profiler: "advanced"
+  # profiler: "pytorch"

conditional-flow-matching/runner/configs/eval.yaml

@@ -0,0 +1,18 @@
+# @package _global_
+
+defaults:
+  - _self_
+  - datamodule: sklearn # choose datamodule with `test_dataloader()` for evaluation
+  - model: cfm
+  - logger: null
+  - trainer: default.yaml
+  - paths: default.yaml
+  - extras: default.yaml
+  - hydra: default.yaml
+
+task_name: "eval"
+
+tags: ["dev"]
+
+# passing checkpoint path is necessary for evaluation
+ckpt_path: ???

conditional-flow-matching/runner/configs/experiment/cfm.yaml

@@ -0,0 +1,22 @@
+# @package _global_
+
+defaults:
+  - override /model: cfm.yaml
+  - override /logger:
+      - csv.yaml
+      - wandb.yaml
+  - override /datamodule: sklearn.yaml
+
+name: "cfm"
+seed: 42
+
+datamodule:
+  batch_size: 512
+
+model:
+  optimizer:
+    weight_decay: 1e-5
+
+trainer:
+  max_epochs: 1000
+  check_val_every_n_epoch: 10

conditional-flow-matching/runner/configs/experiment/cnf.yaml

@@ -0,0 +1,22 @@
+# @package _global_
+
+defaults:
+  - override /model: cnf.yaml
+  - override /logger:
+      - csv.yaml
+      - wandb.yaml
+  - override /datamodule: sklearn.yaml
+
+name: "cnf"
+seed: 42
+
+datamodule:
+  batch_size: 1024
+
+model:
+  optimizer:
+    weight_decay: 1e-5
+
+trainer:
+  max_epochs: 1000
+  check_val_every_n_epoch: 10

conditional-flow-matching/runner/configs/experiment/icnn.yaml

@@ -0,0 +1,18 @@
+# @package _global_
+
+defaults:
+  - override /model: icnn
+  - override /logger:
+      - csv
+      - wandb
+  - override /datamodule: sklearn
+
+name: "icnn"
+seed: 42
+
+datamodule:
+  batch_size: 256
+
+trainer:
+  max_epochs: 10000
+  check_val_every_n_epoch: 100

conditional-flow-matching/runner/configs/experiment/image_cfm.yaml

@@ -0,0 +1,33 @@
+# @package _global_
+
+defaults:
+  - override /model: image_cfm.yaml
+  - override /callbacks: no_stopping
+  - override /logger:
+      - csv.yaml
+      - wandb.yaml
+  - override /datamodule: cifar.yaml
+  - override /trainer: ddp.yaml
+
+name: "cfm"
+seed: 42
+
+datamodule:
+  batch_size: 128
+
+model:
+  _target_: src.models.cfm_module.CFMLitModule
+  sigma_min: 1e-4
+
+  scheduler:
+    _target_: timm.scheduler.PolyLRScheduler
+    _partial_: True
+    warmup_t: 200
+    warmup_lr_init: 1e-8
+    t_initial: 2000
+
+trainer:
+  devices: 2
+  max_epochs: 2000
+  check_val_every_n_epoch: 10
+  limit_val_batches: 0.01

conditional-flow-matching/runner/configs/experiment/image_fm.yaml

@@ -0,0 +1,33 @@
+# @package _global_
+
+defaults:
+  - override /model: image_cfm.yaml
+  - override /callbacks: no_stopping
+  - override /logger:
+      - csv.yaml
+      - wandb.yaml
+  - override /datamodule: cifar.yaml
+  - override /trainer: ddp.yaml
+
+name: "cfm"
+seed: 42
+
+datamodule:
+  batch_size: 128
+
+model:
+  _target_: src.models.cfm_module.FMLitModule
+  sigma_min: 1e-4
+
+  scheduler:
+    _target_: timm.scheduler.PolyLRScheduler
+    _partial_: True
+    warmup_t: 200
+    warmup_lr_init: 1e-8
+    t_initial: 2000
+
+trainer:
+  devices: 2
+  max_epochs: 2000
+  check_val_every_n_epoch: 10
+  limit_val_batches: 0.01

conditional-flow-matching/runner/configs/experiment/image_otcfm.yaml

@@ -0,0 +1,34 @@
+# @package _global_
+
+defaults:
+  - override /model: image_cfm.yaml
+  - override /callbacks: no_stopping
+  - override /logger:
+      - csv.yaml
+      - wandb.yaml
+  - override /datamodule: cifar.yaml
+  - override /trainer: ddp.yaml
+
+name: "cfm"
+seed: 42
+
+datamodule:
+  batch_size: 128
+
+model:
+  _target_: src.models.cfm_module.CFMLitModule
+  sigma_min: 1e-4
+
+  scheduler:
+    _target_: timm.scheduler.PolyLRScheduler
+    _partial_: True
+    warmup_t: 200
+    warmup_lr_init: 1e-8
+    t_initial: 2000
+  ot_sampler: "exact"
+
+trainer:
+  devices: 2
+  max_epochs: 2000
+  check_val_every_n_epoch: 10
+  limit_val_batches: 0.01

conditional-flow-matching/runner/configs/experiment/trajectorynet.yaml

@@ -0,0 +1,22 @@
+# @package _global_
+
+defaults:
+  - override /model: trajectorynet.yaml
+  - override /logger:
+      - csv.yaml
+      - wandb.yaml
+  - override /datamodule: twodim.yaml
+
+name: "cnf"
+seed: 42
+
+datamodule:
+  batch_size: 1024
+
+model:
+  optimizer:
+    weight_decay: 1e-5
+
+trainer:
+  max_epochs: 1000
+  check_val_every_n_epoch: 10

conditional-flow-matching/runner/configs/extras/default.yaml

@@ -0,0 +1,8 @@
+# disable python warnings if they annoy you
+ignore_warnings: False
+
+# ask user for tags if none are provided in the config
+enforce_tags: True
+
+# pretty print config tree at the start of the run using Rich library
+print_config: True

conditional-flow-matching/runner/configs/hparams_search/optuna.yaml

@@ -0,0 +1,49 @@
+# @package _global_
+
+# example hyperparameter optimization of some experiment with Optuna:
+# python train.py -m hparams_search=mnist_optuna experiment=example
+
+defaults:
+  - override /hydra/sweeper: optuna
+
+# choose metric which will be optimized by Optuna
+# make sure this is the correct name of some metric logged in lightning module!
+optimized_metric: "val/2-Wasserstein"
+
+# here we define Optuna hyperparameter search
+# it optimizes for value returned from function with @hydra.main decorator
+# docs: https://hydra.cc/docs/next/plugins/optuna_sweeper
+hydra:
+  mode: "MULTIRUN" # set hydra to multirun by default if this config is attached
+
+  sweeper:
+    _target_: hydra_plugins.hydra_optuna_sweeper.optuna_sweeper.OptunaSweeper
+
+    # storage URL to persist optimization results
+    # for example, you can use SQLite if you set 'sqlite:///example.db'
+    storage: null
+
+    # name of the study to persist optimization results
+    study_name: null
+
+    # number of parallel workers
+    n_jobs: 1
+
+    # 'minimize' or 'maximize' the objective
+    direction: maximize
+
+    # total number of runs that will be executed
+    n_trials: 20
+
+    # choose Optuna hyperparameter sampler
+    # you can choose bayesian sampler (tpe), random search (without optimization), grid sampler, and others
+    # docs: https://optuna.readthedocs.io/en/stable/reference/samplers.html
+    sampler:
+      _target_: optuna.samplers.TPESampler
+      seed: 1234
+      n_startup_trials: 10 # number of random sampling runs before optimization starts
+
+    # define hyperparameter search space
+    params:
+      model.optimizer.lr: interval(0.0001, 0.1)
+      datamodule.batch_size: choice(32, 64, 128, 256)

conditional-flow-matching/runner/configs/hydra/default.yaml

@@ -0,0 +1,15 @@
+# https://hydra.cc/docs/configure_hydra/intro/
+
+# enable color logging
+defaults:
+  - override hydra_logging: colorlog
+  - override job_logging: colorlog
+
+# output directory, generated dynamically on each run
+run:
+  dir: ${paths.log_dir}/${task_name}/runs/${now:%Y-%m-%d}_${now:%H-%M-%S}
+sweep:
+  dir: ${paths.log_dir}/${task_name}/multiruns/${now:%Y-%m-%d}_${now:%H-%M-%S}
+  subdir: ${hydra.job.num}
+job:
+  chdir: true

conditional-flow-matching/runner/configs/launcher/mila_cluster.yaml

@@ -0,0 +1,18 @@
+# @package _global_
+#
+defaults:
+  - override /hydra/launcher: submitit_slurm
+
+hydra:
+  launcher:
+    partition: long
+    cpus_per_task: 2
+    mem_gb: 20
+    gres: gpu:1
+    timeout_min: 1440
+    array_parallelism: 10 # max num of tasks to run in parallel (via job array)
+    setup:
+      - "module purge"
+      - "module load miniconda/3"
+      - "conda activate myenv"
+      - "unset CUDA_VISIBLE_DEVICES"

conditional-flow-matching/runner/configs/launcher/mila_cpu_cluster.yaml

@@ -0,0 +1,16 @@
+# @package _global_
+#
+defaults:
+  - override /hydra/launcher: submitit_slurm
+
+hydra:
+  launcher:
+    partition: long-cpu
+    cpus_per_task: 1
+    mem_gb: 5
+    timeout_min: 100
+    array_parallelism: 64
+    setup:
+      - "module purge"
+      - "module load miniconda/3"
+      - "conda activate myenv"

conditional-flow-matching/runner/configs/local/default.yaml

@@ -0,0 +1,12 @@
+# path to root directory
+# this requires PROJECT_ROOT environment variable to exist
+# PROJECT_ROOT is inferred and set by pyrootutils package in `train.py` and `eval.py`
+root_dir: ${oc.env:PROJECT_ROOT}
+
+scratch_dir: ${oc.env:PROJECT_ROOT}
+
+# path to data directory
+data_dir: ${local.scratch_dir}/data/
+
+# path to logging directory
+log_dir: ${local.scratch_dir}/logs/

conditional-flow-matching/runner/configs/logger/comet.yaml

@@ -0,0 +1,12 @@
+# https://www.comet.ml
+
+comet:
+  _target_: pytorch_lightning.loggers.comet.CometLogger
+  api_key: ${oc.env:COMET_API_TOKEN} # api key is loaded from environment variable
+  save_dir: "${paths.output_dir}"
+  project_name: "lightning-hydra-template"
+  rest_api_key: null
+  # experiment_name: ""
+  experiment_key: null # set to resume experiment
+  offline: False
+  prefix: ""

conditional-flow-matching/runner/configs/logger/csv.yaml

@@ -0,0 +1,7 @@
+# csv logger built in lightning
+
+csv:
+  _target_: pytorch_lightning.loggers.csv_logs.CSVLogger
+  save_dir: "${paths.output_dir}"
+  name: "csv/"
+  prefix: ""

conditional-flow-matching/runner/configs/logger/many_loggers.yaml

@@ -0,0 +1,9 @@
+# train with many loggers at once
+
+defaults:
+  # - comet.yaml
+  - csv.yaml
+  # - mlflow.yaml
+  # - neptune.yaml
+  - tensorboard.yaml
+  - wandb.yaml

conditional-flow-matching/runner/configs/logger/mlflow.yaml

@@ -0,0 +1,12 @@
+# https://mlflow.org
+
+mlflow:
+  _target_: pytorch_lightning.loggers.mlflow.MLFlowLogger
+  # experiment_name: ""
+  # run_name: ""
+  tracking_uri: ${paths.log_dir}/mlflow/mlruns # run `mlflow ui` command inside the `logs/mlflow/` dir to open the UI
+  tags: null
+  # save_dir: "./mlruns"
+  prefix: ""
+  artifact_location: null
+  # run_id: ""

conditional-flow-matching/runner/configs/logger/neptune.yaml

@@ -0,0 +1,9 @@
+# https://neptune.ai
+
+neptune:
+  _target_: pytorch_lightning.loggers.neptune.NeptuneLogger
+  api_key: ${oc.env:NEPTUNE_API_TOKEN} # api key is loaded from environment variable
+  project: username/lightning-hydra-template
+  # name: ""
+  log_model_checkpoints: True
+  prefix: ""

conditional-flow-matching/runner/configs/logger/tensorboard.yaml

@@ -0,0 +1,10 @@
+# https://www.tensorflow.org/tensorboard/
+
+tensorboard:
+  _target_: pytorch_lightning.loggers.tensorboard.TensorBoardLogger
+  save_dir: "${paths.output_dir}/tensorboard/"
+  name: null
+  log_graph: False
+  default_hp_metric: True
+  prefix: ""
+  # version: ""

conditional-flow-matching/runner/configs/logger/wandb.yaml

@@ -0,0 +1,16 @@
+# https://wandb.ai
+
+wandb:
+  _target_: pytorch_lightning.loggers.wandb.WandbLogger
+  # name: "" # name of the run (normally generated by wandb)
+  save_dir: "${paths.output_dir}"
+  offline: False
+  id: null # pass correct id to resume experiment!
+  anonymous: null # enable anonymous logging
+  project: "conditional-flow-model"
+  log_model: False # upload lightning ckpts
+  prefix: "" # a string to put at the beginning of metric keys
+  # entity: "" # set to name of your wandb team
+  group: ""
+  tags: []
+  job_type: ""