squash: merge all unpushed commits

DATASET.md

@@ -0,0 +1,34 @@
+### Dataset
+To download the datataset, run:
+```python
+# download the full dataset
+from huggingface_hub import snapshot_download
+snapshot_download(repo_id="osv5m/osv5m", local_dir="datasets/osv5m", repo_type='dataset')
+```
+
+and finally extract:
+```python
+import os
+import zipfile
+for root, dirs, files in os.walk("datasets/osv5m"):
+    for file in files:
+        if file.endswith(".zip"):
+            with zipfile.ZipFile(os.path.join(root, file), 'r') as zip_ref:
+                zip_ref.extractall(root)
+                os.remove(os.path.join(root, file))
+```
+
+You can also directly load the dataset using `load_dataset`:
+```python
+from datasets import load_dataset
+dataset = load_dataset('osv5m/osv5m', full=False)
+```
+where with `full` you can specify whether you want to load the complete metadata (default: `False`).
+
+If you only want to download the test set, you can run the script below:
+```python
+from huggingface_hub import hf_hub_download
+for i in range(5):
+    hf_hub_download(repo_id="osv5m/osv5m", filename=str(i).zfill(2)+'.zip', subfolder="images/test", repo_type='dataset', local_dir="datasets/osv5m")
+    hf_hub_download(repo_id="osv5m/osv5m", filename="README.md", repo_type='dataset', local_dir="datasets/osv5m")
+```

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Nicolas Dufour
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

callbacks/__init__.py

@@ -0,0 +1,3 @@
+from .ema import EMACallback
+from .fix_nans import FixNANinGrad
+from .data import IncreaseDataEpoch

callbacks/data.py

@@ -0,0 +1,11 @@
+from pytorch_lightning.callbacks import Callback
+
+
+class IncreaseDataEpoch(Callback):
+    def __init__(self):
+        super().__init__()
+
+    def on_train_epoch_start(self, trainer, pl_module):
+        epoch = pl_module.current_epoch
+        if hasattr(trainer.datamodule.train_dataset, "shared_epoch"):
+            trainer.datamodule.train_dataset.shared_epoch.set_value(epoch)

callbacks/ema.py

@@ -0,0 +1,102 @@
+from pytorch_lightning import Callback
+import copy
+import itertools
+import torch
+import contextlib
+from torch.distributed.fsdp import FullyShardedDataParallel
+
+
+class EMACallback(Callback):
+    def __init__(
+        self,
+        module_attr_name,
+        ema_module_attr_name,
+        decay=0.999,
+        start_ema_step=0,
+        init_ema_random=True,
+    ):
+        super().__init__()
+        self.decay = decay
+        self.module_attr_name = module_attr_name
+        self.ema_module_attr_name = ema_module_attr_name
+        self.start_ema_step = start_ema_step
+        self.init_ema_random = init_ema_random
+
+    def on_train_start(self, trainer, pl_module):
+        if pl_module.global_step == 0:
+            if not hasattr(pl_module, self.module_attr_name):
+                raise ValueError(
+                    f"Module {pl_module} does not have attribute {self.module_attr_name}"
+                )
+            if not hasattr(pl_module, self.ema_module_attr_name):
+                pl_module.add_module(
+                    self.ema_module_attr_name,
+                    copy.deepcopy(getattr(pl_module, self.module_attr_name))
+                    .eval()
+                    .requires_grad_(False),
+                )
+            self.reset_ema(pl_module)
+
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
+        if pl_module.global_step == self.start_ema_step:
+            self.reset_ema(pl_module)
+        elif (
+            pl_module.global_step < self.start_ema_step
+            and pl_module.global_step % 100 == 0
+        ):
+            ## slow ema updates for visualisation
+            self.update_ema(pl_module, decay=0.9)
+        elif pl_module.global_step > self.start_ema_step:
+            self.update_ema(pl_module, decay=self.decay)
+
+    def update_ema(self, pl_module, decay=0.999):
+        ema_module = getattr(pl_module, self.ema_module_attr_name)
+        module = getattr(pl_module, self.module_attr_name)
+        context_manager = self.get_model_context_manager(module)
+        with context_manager:
+            with torch.no_grad():
+                ema_params = ema_module.state_dict()
+                for name, param in itertools.chain(
+                    module.named_parameters(), module.named_buffers()
+                ):
+                    if name in ema_params:
+                        if param.requires_grad:
+                            ema_params[name].copy_(
+                                ema_params[name].detach().lerp(param.detach(), decay)
+                            )
+
+    def get_model_context_manager(self, module):
+        fsdp_enabled = is_model_fsdp(module)
+        model_context_manager = contextlib.nullcontext()
+        if fsdp_enabled:
+            model_context_manager = module.summon_full_params(module)
+        return model_context_manager
+
+    def reset_ema(self, pl_module):
+        ema_module = getattr(pl_module, self.ema_module_attr_name)
+        if self.init_ema_random:
+            ema_module.init_weights()
+        else:
+            module = getattr(pl_module, self.module_attr_name)
+            context_manager = self.get_model_context_manager(module)
+            with context_manager:
+                ema_params = ema_module.state_dict()
+                for name, param in itertools.chain(
+                    module.named_parameters(), module.named_buffers()
+                ):
+                    if name in ema_params:
+                        ema_params[name].copy_(param.detach())
+
+
+def is_model_fsdp(model: torch.nn.Module) -> bool:
+    try:
+        if isinstance(model, FullyShardedDataParallel):
+            return True
+
+        # Check if model is wrapped with FSDP
+        for _, obj in model.named_children():
+            if isinstance(obj, FullyShardedDataParallel):
+                return True
+        return False
+    except ImportError:
+        return False

callbacks/fix_nans.py

@@ -0,0 +1,55 @@
+import logging
+from pytorch_lightning.callbacks import Callback
+import torch
+
+log = logging.getLogger(__name__)
+
+
+class FixNANinGrad(Callback):
+    def __init__(self, monitor):
+        super().__init__()
+        self.monitor = monitor
+        self.continuous_nan_batchs = 0
+
+    def on_before_optimizer_step(self, trainer, pl_module, optimizer) -> None:
+        has_nan = []
+        is_inf = []
+        for name, param in pl_module.named_parameters():
+            if param.grad is not None:
+                if torch.isnan(param.grad).any():
+                    has_nan.append(name)
+                if torch.isinf(param.grad).any():
+                    is_inf.append(name)
+                torch.nan_to_num(param.grad, nan=0, posinf=0, neginf=0, out=param.grad)
+        if len(has_nan) > 0:
+            print(f"Found NaN in {has_nan}")
+        if len(is_inf) > 0:
+            print(f"Found Inf in {is_inf}")
+
+    def on_train_batch_end(
+        self,
+        trainer,
+        pl_module,
+        outputs,
+        batch,
+        batch_idx,
+    ) -> None:
+        logs = trainer.callback_metrics
+        i = 0
+        found_metric = False
+        while i < len(self.monitor) and not found_metric:
+            if self.monitor[i] in logs.keys():
+                current = logs[self.monitor[i]].squeeze()
+                found_metric = True
+            else:
+                i += 1
+        if not found_metric:
+            raise ValueError("Asked metric not in logs")
+
+        if not torch.isfinite(current):
+            self.continuous_nan_batchs += 1
+            if self.continuous_nan_batchs >= 5:
+                trainer.should_stop = True
+                log.info("Training interrupted because of NaN in {self.monitor}")
+        else:
+            self.continuous_nan_batchs = 0

configs/computer/a100.yaml

@@ -0,0 +1,8 @@
+devices: 1
+progress_bar_refresh_rate: 2
+num_workers: 8
+sync_batchnorm: False
+accelerator: gpu
+precision: 32
+strategy: auto
+num_nodes: 1

configs/computer/cluster-node-a100.yaml

@@ -0,0 +1,8 @@
+devices: 8
+num_workers: 8
+progress_bar_refresh_rate: 2
+sync_batchnorm: True
+accelerator: gpu
+precision: 32
+strategy: ddp
+num_nodes: 1

configs/computer/cluster-node-v100.yaml

@@ -0,0 +1,8 @@
+devices: 4
+num_workers: 10
+progress_bar_refresh_rate: 2
+sync_batchnorm: True
+accelerator: gpu
+precision: 32
+strategy: ddp
+num_nodes: 1

configs/computer/cpu.yaml

@@ -0,0 +1,8 @@
+devices: null
+num_workers: 0
+progress_bar_refresh_rate: 2
+sync_batchnorm: False
+accelerator: cpu
+precision: 32
+strategy: auto
+num_nodes: null

configs/computer/h100.yaml

@@ -0,0 +1,8 @@
+devices: 1
+progress_bar_refresh_rate: 2
+num_workers: 24
+sync_batchnorm: False
+accelerator: gpu
+precision: 32
+strategy: auto
+num_nodes: 1

configs/computer/v100.yaml

@@ -0,0 +1,8 @@
+devices: 1
+num_workers: 10
+progress_bar_refresh_rate: 2
+sync_batchnorm: False
+accelerator: gpu
+precision: 32
+strategy: auto
+num_nodes: 1

configs/config.yaml

@@ -0,0 +1,90 @@
+defaults:
+  - model: default
+  - computer: v100
+  - dataset: osv5m_emb
+  - stage: null
+  - _self_
+  - exp: ???
+
+model:
+  val_metrics:
+    _target_: metrics.distance_based.HaversineMetrics
+    acc_radiuses:
+      - 1
+      - 25
+      - 200
+      - 750
+      - 2500
+    acc_area: []
+  test_metrics:
+    _target_: metrics.distance_based.HaversineMetrics
+    acc_radiuses:
+      - 1
+      - 25
+      - 200
+      - 750
+      - 2500
+    acc_area: ${areas}
+
+datamodule:
+  _target_: data.datamodule.ImageDataModule
+  train_dataset: ${dataset.train_dataset}
+  val_dataset: ${dataset.val_dataset}
+  test_dataset: ${dataset.test_dataset}
+  full_batch_size: ${dataset.full_batch_size}
+  eval_batch_size: ${dataset.eval_batch_size}
+  num_workers: ${computer.num_workers}
+  num_nodes: ${computer.num_nodes}
+  num_devices: ${computer.devices}
+  val_proportion: 0.02
+
+trainer:
+  _target_: pytorch_lightning.Trainer
+  devices: ${computer.devices}
+  accelerator: ${computer.accelerator}
+  strategy: ${computer.strategy}
+  num_nodes: ${computer.num_nodes}
+  precision: ${computer.precision}
+  max_steps: 1000000
+  val_check_interval: 25000
+  check_val_every_n_epoch: null
+
+logger:
+  _target_: pytorch_lightning.loggers.WandbLogger
+  save_dir: ${root_dir}
+  name: ${experiment_name}${logger_suffix}
+  project: diff_plonk
+  log_model: False
+  offline: False
+
+checkpoints:
+  _target_: pytorch_lightning.callbacks.ModelCheckpoint
+  dirpath: ${root_dir}/checkpoints/${experiment_name}
+  filename: 'epoch_{epoch}'
+  monitor: val/loss
+  save_last: True
+  save_top_k: 0
+  every_n_epochs: 1
+  enable_version_counter: False
+
+progress_bar:
+  _target_: pytorch_lightning.callbacks.TQDMProgressBar
+  refresh_rate: ${computer.progress_bar_refresh_rate}
+
+data_dir: ${root_dir}/datasets
+root_dir:  ${hydra:runtime.cwd}
+experiment_name: ${dataset.name}_${model.name}_${experiment_name_suffix}
+experiment_name_suffix: base
+logger_suffix:  ""
+mode: train # change that to eval to do the testing
+areas: ['country', 'region', 'sub-region', 'city']
+class_name: null
+streetclip: False
+blur: False
+text_tuning: False
+
+hydra:
+  run:
+    dir: outputs/${hydra.job.name}/${now:%Y-%m-%d_%H-%M-%S}/${experiment_name}
+  job:
+    chdir: true

configs/dataset/baselines/im2gps.yaml

@@ -0,0 +1,16 @@
+dataset:
+  name: im2gps
+  full_batch_size: 512
+  test_dataset:
+    _partial_: true
+    _target_: data.data.Baseline
+    path: ${data_dir}/baselines/im2gps
+    which: 'im2gps'
+    transforms: ${dataset.test_transform}
+datamodule:
+  _target_: data.datamodule.BaselineDataModule
+  test_dataset: ${dataset.test_dataset}
+  full_batch_size: ${dataset.full_batch_size}
+  num_workers: ${computer.num_workers}
+  num_nodes: ${computer.num_nodes}
+  num_devices: ${computer.devices}

configs/dataset/baselines/im2gps3k.yaml

@@ -0,0 +1,16 @@
+dataset:
+  name: im2gps3k
+  full_batch_size: 512
+  test_dataset:
+    _partial_: true
+    _target_: data.data.Baseline
+    path: ${data_dir}/baselines/im2gps3k
+    which: 'im2gps3k'
+    transforms: ${dataset.test_transform}
+datamodule:
+  _target_: data.datamodule.BaselineDataModule
+  test_dataset: ${dataset.test_dataset}
+  full_batch_size: ${dataset.full_batch_size}
+  num_workers: ${computer.num_workers}
+  num_nodes: ${computer.num_nodes}
+  num_devices: ${computer.devices}

configs/dataset/baselines/yfcc4k.yaml

@@ -0,0 +1,16 @@
+dataset:
+  name: yfcc4k
+  full_batch_size: 512
+  test_dataset:
+    _partial_: true
+    _target_: data.data.Baseline
+    path: ${data_dir}/baselines/yfcc4k
+    which: 'yfcc4k'
+    transforms: ${dataset.test_transform}
+datamodule:
+  _target_: data.datamodule.BaselineDataModule
+  test_dataset: ${dataset.test_dataset}
+  full_batch_size: ${dataset.full_batch_size}
+  num_workers: ${computer.num_workers}
+  num_nodes: ${computer.num_nodes}
+  num_devices: ${computer.devices}

configs/dataset/combined_emb.yaml

@@ -0,0 +1,38 @@
+defaults:
+  - train_transform: empty
+  - test_transform: empty
+  - _self_
+
+name: iNaturalist_OSV5M_YFCC100M_${dataset.embedding_name}
+full_batch_size: 2048
+cond_dim: 1024
+eval_batch_size: 4096
+output_type: emb
+embedding_name: dinov2_vitl14_registers
+
+train_dataset:
+  _partial_: true
+  _target_: data.webdataset.GPSWebdataset
+  root: ${data_dir}/YFCC100M/train/ ${data_dir}/osv5m/train/ ${data_dir}/inaturalist/train/ ${data_dir}/osv5m/train/ ${data_dir}/inaturalist/train/
+  train: true
+  embedding_name: ${dataset.embedding_name}
+  return_image: false
+  metadata_attributes: []
+
+val_dataset:
+  _partial_: true
+  _target_: data.webdataset.GPSWebdataset
+  root: ${data_dir}/YFCC100M/yfcc4k/
+  train: false
+  embedding_name: ${dataset.embedding_name}
+  return_image: false
+  metadata_attributes: []
+
+test_dataset:
+  _partial_: true
+  _target_: data.webdataset.GPSWebdataset
+  root: ${data_dir}/YFCC100M/yfcc4k/
+  train: false
+  embedding_name: ${dataset.embedding_name}
+  return_image: false
+  metadata_attributes: []

configs/dataset/inaturalist_emb.yaml

@@ -0,0 +1,38 @@
+defaults:
+  - train_transform: empty
+  - test_transform: empty
+  - _self_
+
+name: iNaturalist_${dataset.embedding_name}
+full_batch_size: 512
+cond_dim: 1024
+eval_batch_size: 4096
+output_type: emb
+embedding_name: dinov2_vitl14_registers
+
+train_dataset:
+  _partial_: true
+  _target_: data.webdataset.GPSWebdataset
+  root: ${data_dir}/inaturalist/train/
+  train: true
+  embedding_name: ${dataset.embedding_name}
+  return_image: false
+  metadata_attributes: []
+
+val_dataset:
+  _partial_: true
+  _target_: data.webdataset.GPSWebdataset
+  root: ${data_dir}/inaturalist/val/
+  train: false
+  embedding_name: ${dataset.embedding_name}
+  return_image: false
+  metadata_attributes: []
+
+test_dataset:
+  _partial_: true
+  _target_: data.webdataset.GPSWebdataset
+  root: ${data_dir}/inaturalist/test/
+  train: false
+  embedding_name: ${dataset.embedding_name}
+  return_image: false
+  metadata_attributes: []

configs/dataset/osv5m.yaml

@@ -0,0 +1,43 @@
+defaults:
+  - train_transform: fast_clip
+  - test_transform: fast_clip
+  - _self_
+
+name: osv5m
+full_batch_size: 2048
+eval_batch_size: 4096
+train_dataset:
+  _partial_: true
+  _target_: data.data.OSV5M
+  path: ${data_dir}/osv5m/
+  split: train
+  class_name: ${class_name}
+  transforms: ${dataset.train_transform}
+  is_baseline: ${is_baseline}
+  areas: ${areas}
+  streetclip: ${streetclip}
+  blur: ${blur}
+
+val_dataset:
+  _partial_: true
+  _target_: data.data.OSV5M
+  path: ${data_dir}/osv5m/
+  split: val
+  class_name: ${class_name}
+  transforms: ${dataset.test_transform}
+  is_baseline: ${is_baseline}
+  areas: ${areas}
+  streetclip: ${streetclip}
+  blur: ${blur}
+
+test_dataset:
+  _partial_: true
+  _target_: data.data.OSV5M
+  path: ${data_dir}/osv5m/
+  split: test
+  class_name: ${class_name}
+  transforms: ${dataset.test_transform}
+  is_baseline: ${is_baseline}
+  areas: ${areas}
+  streetclip: ${streetclip}
+  blur: ${blur}

configs/dataset/osv5m_emb.yaml

@@ -0,0 +1,38 @@
+defaults:
+  - train_transform: empty
+  - test_transform: empty
+  - _self_
+
+name: osv5m_${dataset.embedding_name}
+full_batch_size: 1024
+eval_batch_size: 4096
+cond_dim: 1024
+output_type: emb
+embedding_name: street_clip
+
+train_dataset:
+  _partial_: true
+  _target_: data.webdataset.GPSWebdataset
+  root: ${data_dir}/osv5m/train/
+  train: true
+  embedding_name: ${dataset.embedding_name}
+  return_image: false
+  metadata_attributes: []
+
+val_dataset:
+  _partial_: true
+  _target_: data.webdataset.GPSWebdataset
+  root: ${data_dir}/osv5m/val/
+  train: false
+  embedding_name: ${dataset.embedding_name}
+  return_image: false
+  metadata_attributes: ["unique_country", "unique_region", "unique_sub-region", "unique_city"]
+
+test_dataset:
+  _partial_: true
+  _target_: data.webdataset.GPSWebdataset
+  root: ${data_dir}/osv5m/test/
+  train: false
+  embedding_name: ${dataset.embedding_name}
+  return_image: false
+  metadata_attributes: ["unique_country", "unique_region", "unique_sub-region", "unique_city"]

configs/dataset/test_transform/center_crop.yaml

@@ -0,0 +1,12 @@
+_target_: torchvision.transforms.Compose
+transforms:
+  - _target_: torchvision.transforms.ToTensor
+  - _target_: utils.image_processing.CenterCrop
+    ratio: "1:1"
+  - _target_: torchvision.transforms.Resize
+    size: ${dataset.img_resolution}
+    interpolation: 3
+    antialias: true
+  - _target_: torchvision.transforms.Normalize
+    mean: 0.5
+    std: 0.5

configs/dataset/test_transform/clip.yaml

@@ -0,0 +1,2 @@
+_target_: data.transforms.ClipTransform
+split: val

configs/dataset/test_transform/empty.yaml

@@ -0,0 +1,2 @@
+_target_: data.data.null_transform
+_partial_: true

configs/dataset/test_transform/fast_clip.yaml

@@ -0,0 +1,12 @@
+_target_: torchvision.transforms.Compose
+transforms:
+  - _target_: torchvision.transforms.Resize
+    size: 224
+    interpolation: 3
+    antialias: true
+  - _target_: torchvision.transforms.CenterCrop
+    size: 224
+  - _target_: torchvision.transforms.ToTensor
+  - _target_: torchvision.transforms.Normalize
+    mean: [0.48145466, 0.4578275, 0.40821073]
+    std: [0.26862954, 0.26130258, 0.27577711]

configs/dataset/test_transform/fast_resnet.yaml

@@ -0,0 +1,12 @@
+_target_: torchvision.transforms.Compose
+transforms:
+  - _target_: torchvision.transforms.Resize
+    size: 224
+    interpolation: 3
+    antialias: true
+  - _target_: torchvision.transforms.CenterCrop
+    size: 224
+  - _target_: torchvision.transforms.ToTensor
+  - _target_: torchvision.transforms.Normalize
+    mean: [0.485 ,0.456 ,0.406]
+    std: [0.229, 0.224, 0.225]

configs/dataset/test_transform/none.yaml

@@ -0,0 +1,6 @@
+_target_: torchvision.transforms.Compose
+transforms:
+  - _target_: torchvision.transforms.ToTensor
+  - _target_: torchvision.transforms.Normalize
+    mean: 0.5
+    std: 0.5

configs/dataset/train_transform/augmentation.yaml

@@ -0,0 +1,85 @@
+_target_: data.augmentation.ImageAugmentation
+names: "standard_augmentation,geometric_augmentation,clip_transform"
+
+# always apply clip_transform at the end
+clip_transform:
+  _target_: torchvision.transforms.Compose
+  transforms:
+    - _target_: torchvision.transforms.Resize
+      size: 224
+      interpolation: 3
+      antialias: true
+    - _target_: torchvision.transforms.CenterCrop
+      size: 224
+    - _target_: torchvision.transforms.ToTensor
+    - _target_: torchvision.transforms.Normalize
+      mean: [0.48145466, 0.4578275, 0.40821073]
+      std: [0.26862954, 0.26130258, 0.27577711]
+
+standard_augmentation:
+  _target_: data.augmentation.StandardAugmentation
+  # by default, we all augmentation methods
+  names: "brightness,contrast,sharpness,color,blur,gaussian_noise"
+
+  # random PIL brigtness
+  brightness:
+    _target_: data.augmentation.PillowBrightness
+    p: 0.2
+    factor_interval: [0.5, 1.5]
+
+  # random PIL contrast
+  contrast:
+    _target_: data.augmentation.PillowContrast
+    p: 0.2
+    factor_interval: [0.3, 3]
+
+  # random PIL sharpness
+  sharpness:
+    _target_: data.augmentation.PillowSharpness
+    p: 0.2
+    factor_interval: [0.5, 30.0]
+
+  # random PIL color
+  color:
+    _target_: data.augmentation.PillowColor
+    p: 0.2
+    factor_interval: [0.0, 2.0]
+
+  # random PIL blur
+  blur:
+    _target_: data.augmentation.PillowBlur
+    p: 0.2
+    factor_interval: [1, 2]
+
+  # random numpy gaussian noise
+  gaussian_noise:
+    _target_: data.augmentation.NumpyGaussianNoise
+    p: 0.2
+    factor_interval: [0.1, 0.04]
+
+geometric_augmentation:
+  _target_: data.augmentation.GeometricAugmentation
+  # by default, we all augmentation methods
+  names: "random_rotation,random_resized_crop,random_horizontal_flip"
+
+  # random rotation
+  random_rotation:
+    _target_: torchvision.transforms.RandomRotation
+    degrees: [-15, 15]
+
+  # random crop
+  random_resized_crop:
+    _target_: torchvision.transforms.RandomResizedCrop
+    scale: [0.5, 1.0]
+    ratio: [0.9, 1.1]
+    size: 224
+
+  # random horizontal flip
+  random_horizontal_flip:
+    _target_: torchvision.transforms.RandomHorizontalFlip
+    p: 0.5
+
+  # random vertical flip
+  random_vertical_flip:
+    _target_: torchvision.transforms.RandomVerticalFlip
+    p: 0.5

configs/dataset/train_transform/center_crop.yaml

@@ -0,0 +1,14 @@
+_target_: torchvision.transforms.Compose
+transforms:
+  - _target_: torchvision.transforms.ToTensor
+  - _target_: utils.image_processing.CenterCrop
+    ratio: "1:1"
+  - _target_: torchvision.transforms.Resize
+    size: ${dataset.img_resolution}
+    interpolation: 3
+    antialias: true
+  - _target_: torchvision.transforms.RandomHorizontalFlip
+    p: 0.5
+  - _target_: torchvision.transforms.Normalize
+    mean: 0.5
+    std: 0.5

configs/dataset/train_transform/clip.yaml

@@ -0,0 +1,2 @@
+_target_: data.transforms.ClipTransform
+split: val

configs/dataset/train_transform/empty.yaml

@@ -0,0 +1,2 @@
+_target_: data.data.null_transform
+_partial_: true

configs/dataset/train_transform/fast_clip.yaml

@@ -0,0 +1,12 @@
+_target_: torchvision.transforms.Compose
+transforms:
+  - _target_: torchvision.transforms.Resize
+    size: 224
+    interpolation: 3
+    antialias: true
+  - _target_: torchvision.transforms.CenterCrop
+    size: 224
+  - _target_: torchvision.transforms.ToTensor
+  - _target_: torchvision.transforms.Normalize
+    mean: [0.48145466, 0.4578275, 0.40821073]
+    std: [0.26862954, 0.26130258, 0.27577711]

configs/dataset/train_transform/fast_resnet.yaml

@@ -0,0 +1,12 @@
+_target_: torchvision.transforms.Compose
+transforms:
+  - _target_: torchvision.transforms.Resize
+    size: 224
+    interpolation: 3
+    antialias: true
+  - _target_: torchvision.transforms.CenterCrop
+    size: 224
+  - _target_: torchvision.transforms.ToTensor
+  - _target_: torchvision.transforms.Normalize
+    mean: [0.485 ,0.456 ,0.406]
+    std: [0.229, 0.224, 0.225]

configs/dataset/train_transform/none.yaml

@@ -0,0 +1,7 @@
+_target_: torchvision.transforms.Compose
+transforms:
+  - _target_: torchvision.transforms.Resize
+    size: 224
+    interpolation: 3
+    antialias: true
+  - _target_: torchvision.transforms.ToTensor

configs/dataset/yfcc_emb.yaml

@@ -0,0 +1,38 @@
+defaults:
+  - train_transform: empty
+  - test_transform: empty
+  - _self_
+
+name: iNaturalist_${dataset.embedding_name}
+full_batch_size: 2048
+cond_dim: 1024
+eval_batch_size: 4096
+output_type: emb
+embedding_name: dinov2_vitl14_registers
+
+train_dataset:
+  _partial_: true
+  _target_: data.webdataset.GPSWebdataset
+  root: ${data_dir}/YFCC100M/train/
+  train: true
+  embedding_name: ${dataset.embedding_name}
+  return_image: false
+  metadata_attributes: []
+
+val_dataset:
+  _partial_: true
+  _target_: data.webdataset.GPSWebdataset
+  root: ${data_dir}/YFCC100M/yfcc4k/
+  train: false
+  embedding_name: ${dataset.embedding_name}
+  return_image: false
+  metadata_attributes: []
+
+test_dataset:
+  _partial_: true
+  _target_: data.webdataset.GPSWebdataset
+  root: ${data_dir}/YFCC100M/yfcc4k/
+  train: false
+  embedding_name: ${dataset.embedding_name}
+  return_image: false
+  metadata_attributes: []

configs/exp/YFCC100M_geoadalnmlp_r2_small_sigmoid_diffusion.yaml

@@ -0,0 +1,35 @@
+# @package _global_
+
+defaults:
+  - override /dataset: yfcc_emb
+  - override /model: emb_cond
+  - override /model/network: geo_adaln_mlp
+  - override /model/train_noise_scheduler: sigmoid
+  - override /model/inference_noise_scheduler: sigmoid
+  - override /model/loss: ddpm
+  - _self_
+
+model:
+  network:
+    depth: 12
+    dim: 512
+  optimizer:
+    optim:
+      lr: 8e-4
+      weight_decay: 0.05
+  loss:
+    cond_drop_rate: 0.1
+  train_noise_scheduler:
+    start: -7
+    end: 3
+    tau: 1.0
+  inference_noise_scheduler:
+    start: -7
+    end: 3
+    tau: 1.0
+  interpolant: diffusion
+dataset:
+  full_batch_size: 1024
+
+experiment_name_suffix: small_sigmoid
+areas: []

configs/exp/YFCC100M_geoadalnmlp_r3_small_linear_flow_rieman.yaml

@@ -0,0 +1,32 @@
+# @package _global_
+
+defaults:
+  - override /dataset: yfcc_emb
+  - override /model: emb_cond_cartesian
+  - override /model/network: geo_adaln_mlp
+  - override /model/train_noise_scheduler: linear
+  - override /model/inference_noise_scheduler: linear
+  - override /model/loss: riemannian_flow_matching
+  - override /model/manifold: sphere
+  - override /model/val_sampler: riemannian_flow_matching
+  - override /model/test_sampler: riemannian_flow_matching
+  - _self_
+
+model:
+  network:
+    depth: 12
+    dim: 512
+  optimizer:
+    optim:
+      lr: 8e-4
+      weight_decay: 0.05
+  loss:
+    cond_drop_rate: 0.1
+  interpolant: flow_matching
+
+dataset:
+  full_batch_size: 1024
+
+areas: []
+
+experiment_name_suffix: small_sigmoid

configs/exp/YFCC100M_geoadalnmlp_r3_small_sigmoid_diffusion.yaml

@@ -0,0 +1,36 @@
+# @package _global_
+
+defaults:
+  - override /dataset: yfcc_emb
+  - override /model: emb_cond_cartesian
+  - override /model/network: geo_adaln_mlp
+  - override /model/train_noise_scheduler: sigmoid
+  - override /model/inference_noise_scheduler: sigmoid
+  - override /model/loss: ddpm
+  - _self_
+
+model:
+  network:
+    depth: 12
+    dim: 512
+  optimizer:
+    optim:
+      lr: 8e-4
+      weight_decay: 0.05
+  loss:
+    cond_drop_rate: 0.1
+  train_noise_scheduler:
+    start: -7
+    end: 3
+    tau: 1.0
+  inference_noise_scheduler:
+    start: -7
+    end: 3
+    tau: 1.0
+  interpolant: diffusion
+
+dataset:
+  full_batch_size: 1024
+
+experiment_name_suffix: small_sigmoid
+areas: []

configs/exp/YFCC100M_geoadalnmlp_r3_small_sigmoid_flow.yaml

@@ -0,0 +1,38 @@
+# @package _global_
+
+defaults:
+  - override /dataset: yfcc_emb
+  - override /model: emb_cond_cartesian
+  - override /model/network: geo_adaln_mlp
+  - override /model/train_noise_scheduler: sigmoid
+  - override /model/inference_noise_scheduler: sigmoid
+  - override /model/loss: flow_matching
+  - override /model/val_sampler: flow_matching
+  - override /model/test_sampler: flow_matching
+  - _self_
+
+model:
+  network:
+    depth: 12
+    dim: 512
+  optimizer:
+    optim:
+      lr: 8e-4
+      weight_decay: 0.05
+  loss:
+    cond_drop_rate: 0.1
+  train_noise_scheduler:
+    start: -7
+    end: 3
+    tau: 1.0
+  inference_noise_scheduler:
+    start: -7
+    end: 3
+    tau: 1.0
+  interpolant: flow_matching
+
+dataset:
+  full_batch_size: 1024
+
+experiment_name_suffix: small_sigmoid
+areas: []

configs/exp/YFCC100M_geoadalnmlp_r3_small_sigmoid_flow_riemann.yaml

@@ -0,0 +1,40 @@
+# @package _global_
+
+defaults:
+  - override /dataset: yfcc_emb
+  - override /model: emb_cond_cartesian
+  - override /model/network: geo_adaln_mlp
+  - override /model/train_noise_scheduler: sigmoid
+  - override /model/inference_noise_scheduler: sigmoid
+  - override /model/loss: riemannian_flow_matching
+  - override /model/manifold: sphere
+  - override /model/val_sampler: riemannian_flow_matching
+  - override /model/test_sampler: riemannian_flow_matching
+  - _self_
+
+model:
+  network:
+    depth: 12
+    dim: 512
+  optimizer:
+    optim:
+      lr: 8e-4
+      weight_decay: 0.05
+  loss:
+    cond_drop_rate: 0.1
+  train_noise_scheduler:
+    start: -7
+    end: 3
+    tau: 1.0
+  inference_noise_scheduler:
+    start: -7
+    end: 3
+    tau: 1.0
+  interpolant: flow_matching
+
+dataset:
+  full_batch_size: 1024
+
+areas: []
+
+experiment_name_suffix: small_sigmoid

configs/exp/YFCC100M_geoadalnmlp_von_fisher.yaml

@@ -0,0 +1,26 @@
+# @package _global_
+
+defaults:
+  - override /dataset: yfcc_emb
+  - override /model: von_fisher
+  - override /model/network: geo_adaln_mlp_von_fisher
+  - override /model/loss: von_fisher
+  - override /model/val_sampler: von_fisher
+  - override /model/test_sampler: von_fisher
+  - _self_
+
+model:
+  network:
+    depth: 11 # To compensate the increase in params
+    dim: 512
+  optimizer:
+    optim:
+      lr: 1e-4
+      weight_decay: 0.05
+dataset:
+  full_batch_size: 1024
+trainer:
+  gradient_clip_val: 0.05
+  gradient_clip_algorithm: norm
+areas: []
+experiment_name_suffix: von_fisher

configs/exp/YFCC100M_geoadalnmlp_von_fisher_mixture.yaml

@@ -0,0 +1,26 @@
+# @package _global_
+
+defaults:
+  - override /dataset: yfcc_emb
+  - override /model: von_fisher_mixture
+  - override /model/network: geo_adaln_mlp_von_fisher_mixture
+  - override /model/loss: von_fisher_mixture
+  - override /model/val_sampler: von_fisher_mixture
+  - override /model/test_sampler: von_fisher_mixture
+  - _self_
+
+model:
+  network:
+    depth: 11 # To compensate the increase in params
+    dim: 512
+  optimizer:
+    optim:
+      lr: 1e-5
+      weight_decay: 0.05
+dataset:
+  full_batch_size: 1024
+trainer:
+  gradient_clip_val: 0.01
+  gradient_clip_algorithm: norm
+experiment_name_suffix: von_fisher_mixture
+areas: []

configs/exp/combined_geoadalnmlp_r3_small_sigmoid_flow_riemann.yaml

@@ -0,0 +1,40 @@
+# @package _global_
+
+defaults:
+  - override /dataset: combined_emb
+  - override /model: emb_cond_cartesian
+  - override /model/network: geo_adaln_mlp
+  - override /model/train_noise_scheduler: sigmoid
+  - override /model/inference_noise_scheduler: sigmoid
+  - override /model/loss: riemannian_flow_matching
+  - override /model/manifold: sphere
+  - override /model/val_sampler: riemannian_flow_matching
+  - override /model/test_sampler: riemannian_flow_matching
+  - _self_
+
+model:
+  network:
+    depth: 12
+    dim: 512
+  optimizer:
+    optim:
+      lr: 8e-4
+      weight_decay: 0.05
+  loss:
+    cond_drop_rate: 0.1
+  train_noise_scheduler:
+    start: -7
+    end: 3
+    tau: 1.0
+  inference_noise_scheduler:
+    start: -7
+    end: 3
+    tau: 1.0
+  interpolant: flow_matching
+
+dataset:
+  full_batch_size: 1024
+
+areas: []
+
+experiment_name_suffix: small_sigmoid

configs/exp/iNaturalist_geoadalnmlp_r2_small_sigmoid_diffusion.yaml

@@ -0,0 +1,36 @@
+# @package _global_
+
+defaults:
+  - override /dataset: inaturalist_emb
+  - override /model: emb_cond
+  - override /model/network: geo_adaln_mlp
+  - override /model/train_noise_scheduler: sigmoid
+  - override /model/inference_noise_scheduler: sigmoid
+  - override /model/loss: ddpm
+  - _self_
+
+model:
+  network:
+    depth: 12
+    dim: 256
+  optimizer:
+    optim:
+      lr: 8e-4
+      weight_decay: 0.1
+  loss:
+    cond_drop_rate: 0.1
+  train_noise_scheduler:
+    start: -7
+    end: 3
+    tau: 1.0
+  inference_noise_scheduler:
+    start: -7
+    end: 3
+    tau: 1.0
+  interpolant: diffusion
+dataset:
+  full_batch_size: 512
+
+areas: []
+
+experiment_name_suffix: small_sigmoid

configs/exp/iNaturalist_geoadalnmlp_r3_small_sigmoid_diffusion.yaml

@@ -0,0 +1,37 @@
+# @package _global_
+
+defaults:
+  - override /dataset: inaturalist_emb
+  - override /model: emb_cond_cartesian
+  - override /model/network: geo_adaln_mlp
+  - override /model/train_noise_scheduler: sigmoid
+  - override /model/inference_noise_scheduler: sigmoid
+  - override /model/loss: ddpm
+  - _self_
+
+model:
+  network:
+    depth: 12
+    dim: 256
+  optimizer:
+    optim:
+      lr: 8e-4
+      weight_decay: 0.1
+  loss:
+    cond_drop_rate: 0.1
+  train_noise_scheduler:
+    start: -7
+    end: 3
+    tau: 1.0
+  inference_noise_scheduler:
+    start: -7
+    end: 3
+    tau: 1.0
+  interpolant: diffusion
+
+dataset:
+  full_batch_size: 512
+
+areas: []
+
+experiment_name_suffix: small_sigmoid

configs/exp/iNaturalist_geoadalnmlp_r3_small_sigmoid_flow.yaml

@@ -0,0 +1,39 @@
+# @package _global_
+
+defaults:
+  - override /dataset: inaturalist_emb
+  - override /model: emb_cond_cartesian
+  - override /model/network: geo_adaln_mlp
+  - override /model/train_noise_scheduler: sigmoid
+  - override /model/inference_noise_scheduler: sigmoid
+  - override /model/loss: flow_matching
+  - override /model/val_sampler: flow_matching
+  - override /model/test_sampler: flow_matching
+  - _self_
+
+model:
+  network:
+    depth: 12
+    dim: 256
+  optimizer:
+    optim:
+      lr: 8e-4
+      weight_decay: 0.1
+  loss:
+    cond_drop_rate: 0.1
+  train_noise_scheduler:
+    start: -7
+    end: 3
+    tau: 1.0
+  inference_noise_scheduler:
+    start: -7
+    end: 3
+    tau: 1.0
+  interpolant: flow_matching
+
+dataset:
+  full_batch_size: 512
+
+areas: []
+
+experiment_name_suffix: small_sigmoid