Merge pull request #311 from alan-turing-institute/run_diffusion_models

Add diffusion model configs

Author: radka-j

local_hydra/local_experiment/epd_diffusion_dm_256_dc_large.yaml.yaml

@@ -0,0 +1,34 @@
+# @package _global_
+defaults:
+  - /distributed: ddp_4gpu_slurm
+  - override /datamodule: advection_diffusion_multichannel_64_64
+  - override /encoder@model.encoder: dc_large
+  - override /decoder@model.decoder: dc_large
+  - override /processor@model.processor: diffusion_vit
+  - _self_
+
+experiment_name: epd_diffusion_dm_256
+
+datamodule:
+  use_normalization: true
+  batch_size: 256
+
+logging:
+  wandb:
+    enabled: true
+
+optimizer:
+  learning_rate: 0.0002
+
+model:
+  encoder:
+    periodic: false
+  decoder:
+    periodic: false
+  train_in_latent_space: true
+  processor:
+    backbone:
+      hid_channels: 640
+
+trainer:
+  max_time: "00:12:00:00"

local_hydra/local_experiment/epd_diffusion_dm_256_dc_large_periodic.yaml

@@ -0,0 +1,34 @@
+# @package _global_
+defaults:
+  - /distributed: ddp_4gpu_slurm
+  - override /datamodule: advection_diffusion_multichannel_64_64
+  - override /encoder@model.encoder: dc_large
+  - override /decoder@model.decoder: dc_large
+  - override /processor@model.processor: diffusion_vit
+  - _self_
+
+experiment_name: epd_diffusion_dm_256
+
+datamodule:
+  use_normalization: true
+  batch_size: 256
+
+logging:
+  wandb:
+    enabled: true
+
+optimizer:
+  learning_rate: 0.0002
+
+model:
+  encoder:
+    periodic: true
+  decoder:
+    periodic: true
+  train_in_latent_space: true
+  processor:
+    backbone:
+      hid_channels: 640
+
+trainer:
+  max_time: "00:12:00:00"