Diffusion models

examples/basic_generate.py

@@ -0,0 +1,48 @@
+import argparse
+import os
+import torch
+from ruamel.yaml import YAML
+from ruamel.yaml.comments import CommentedMap as ruamelDict
+import argparse
+import os
+import torch
+from ruamel.yaml import YAML
+from ruamel.yaml.comments import CommentedMap as ruamelDict
+from matey import Generator
+from matey.utils import setup_dist, check_sp, profile_function, log_to_file, log_versions, YParams
+import glob, socket
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_dir", default='./Demo_Diffusion_CIFAR10_3level_B_1e-3lr/basic_config/demo_diffusion/', type=str)
+    parser.add_argument("--yaml_config", default='hyperparams.yaml', type=str)
+    parser.add_argument("--use_ddp", action='store_true', help='Use distributed data parallel')
+    parser.add_argument("--config", default='basic_config', type=str)
+    parser.add_argument("--output_dir", default='./CIFAR10_generation_outputs/', type=str)
+
+    args = parser.parse_args()
+    params = YParams(os.path.join(args.model_dir, args.yaml_config))
+    params.use_ddp = args.use_ddp
+    params['output_dir'] = args.output_dir
+
+    os.makedirs(params.output_dir, exist_ok=True)
+
+    # Set up distributed training
+    device, world_size, local_rank, global_rank = setup_dist(params)
+    print(f"local_rank={local_rank}, global_rank={global_rank}, world_size={world_size}, host={socket.gethostname()}", flush=True)
+
+    # Modify params
+    params['batch_size'] =int(params.batch_size//world_size)
+    params['checkpoint_path'] = os.path.join(args.model_dir, 'training_checkpoints/ckpt.tar')
+    params['best_checkpoint_path'] = os.path.join(args.model_dir, 'training_checkpoints/best_ckpt.tar')
+
+    assert os.path.isfile(params.checkpoint_path), f"file {params.checkpoint_path} not found" 
+    assert os.path.isfile(params.best_checkpoint_path), f"file {params.best_checkpoint_path} not found" 
+    params['resuming'] = True 
+
+    generator = Generator(params, global_rank, local_rank, device)
+
+    generator.generate(seed=42, num_samples=9, batch_list=[6])
+
+    if params.log_to_screen:
+        print('DONE ---- rank %d'%global_rank)

examples/config/Demo_MW_diffusion_TT.yaml

@@ -0,0 +1,91 @@
+basic_config: &basic_config
+  # Run settings
+  log_to_screen: !!bool True # Log progress to screen.
+  save_checkpoint: !!bool True # Save checkpoints
+  checkpoint_save_interval: 10 # Save every # epochs - also saves "best" according to val loss
+  true_time: !!bool False # Debugging setting - sets num workers to zero and activates syncs
+  num_data_workers: 1 # Generally pulling 8 cpu per process, so using 6 for DL - not sure if best ratio
+  enable_amp: !!bool False # Use automatic mixed precision - blows up with low variance fields right now
+  compile: !!bool False # Compile model - Does not currently work
+  gradient_checkpointing: !!bool False # Whether to use gradient checkpointing - Slow, but lower memory
+  exp_dir: 'Demo_Diffusion_MW_cond_TT_3level_S_lt5' # Output path
+  # exp_dir: 'Demo_Diffusion_CIFAR10_correctedEpoch' # Output path
+  log_interval: 1 # How often to log - Don't think this is actually implemented
+  pretrained: !!bool False # Whether to load a pretrained model
+  # Training settings
+  drop_path: 0.1
+  batch_size: 32 #1
+  max_epochs: 200 #800
+  scheduler_epochs: -1
+  epoch_size: 100 #100 #2000 # Artificial epoch size
+  rescale_gradients: !!bool False # Activate hook that scales block gradients to norm 1
+  optimizer: 'AdamW' # a
+  scheduler: 'cosine' # Only cosine implemented
+  warmup_steps: 10 # Warmup when not using DAdapt
+  learning_rate: 1e-3 # -1 means use DAdapt
+  weight_decay: 1e-3 
+  n_states: 15 #12   # Number of state variables across the datasets - Can be larger than real number and things will just go unused
+  state_names: ['Pressure', 'Vx', 'Vy', 'Density',  'Vx', 'Vy', 'Density', 'Pressure'] # Should be sorted
+  dt: 1 # Striding of data - Not currently implemented > 1
+  leadtime_max: 5 #prediction lead time range [1, leadtime_max]
+  n_steps: 1 # Length of history to include in input
+  enforce_max_steps: !!bool False # If false and n_steps > dataset steps, use dataset steps. Otherwise, raise Exception.
+  accum_grad: 1  # Real batch size is accum * batch_size, real steps/"epoch" is epoch_size / accum
+  # Model settings
+  model_type: 'turbt'
+  # model_type:  'vit_all2all' # no need for time_type and space_type inputs
+  #model_type:  'svit' #currently only support time_type=="all2all_time" and space_type=="all2all"
+  #time_type: 'all2all_time' # 
+  #space_type: 'all2all' #
+  #model_type:  'avit' #currently only support space_type=="axial_attention" and time_type=="attention"
+  #time_type: 'attention' # Conditional on block type
+  #space_type: 'axial_attention' # Conditional on block type 
+  tie_fields: !!bool False # Whether to use 1 embedding per field per data
+  embed_dim: 384 # Dimension of internal representation - 192/384/768/1024 for Ti/S/B/L
+  num_heads: 6 # Number of heads for attention - 3/6/12/16 for Ti/S/B/L
+  processor_blocks: 12 # Number of transformer blocks in the backbone - 12/12/12/24 for Ti/S/B/L
+  diffusion_config:
+    diffusion: !!bool True
+    cond_diffusion: !!bool True
+    model_channels: 128       # Base channel count for noise embedding MLP
+    channel_mult_emb: 4       # Multiplier for embedding MLP hidden dim (emb_channels = model_channels * channel_mult_emb)
+    embedding_type: 'positional'  # Noise level embedding type: 'positional' or 'fourier'
+    channel_mult_noise: 1     # Multiplier for noise embedding dim (noise_channels = model_channels * channel_mult_noise)
+  tokenizer_heads:
+    - head_name: "tk-2D"
+      patch_size: [[1, 4, 4]]
+  hierarchical:
+      filtersize: 2
+      nlevels: 3 #[2^3, 2^2, 2^1, 2^0]
+  adaptive: !!bool False
+  sts_model: !!bool False
+  adap_samp: !!bool False
+  nrefines: 0 #number of coarse tokens picked to be refined
+  bias_type: 'none'  # Options rel, continuous, none
+  # Data settings
+  train_val_test: [.8, .1, .1]
+  augmentation: !!bool False # Augmentation not implemented
+  use_all_fields: !!bool True # Prepopulate the field metadata dictionary from dictionary in datasets
+  tie_batches: !!bool False # Force everything in batch to come from one dset
+  extended_names: !!bool False # Whether to use extended names - not currently implemented
+  embedding_offset: 0  # Use when adding extra finetuning fields
+  train_data_paths: [
+              #['/home/6pz/data/PDEBench_Data/2D/shallow-water', 'swe', ''],
+              #['/home/6pz/data/PDEBench_Data/2D/NS_incom', 'incompNS', ''],
+              #['/home/6pz/data/PDEBench_Data/2D/CFD/2D_Train_Rand', compNS, '128'],
+              #['/home/6pz/data/PDEBench_Data/2D/CFD/2D_Train_Rand', compNS, '512'],
+              #['/home/6pz/data/PDEBench_Data/2D/CFD/2D_Train_Turb', compNS, ''],
+              #['/home/6pz/data/PDEBench_Data/2D/diffusion-reaction', 'diffre2d', ''],
+              ['/lustre/orion/lrn037/proj-shared/miniweather_thermals_tiny/train', 'thermalcollision2d', '', 'tk-2D'],              
+              ]
+  valid_data_paths: [
+              #['/home/6pz/data/PDEBench_Data/2D/shallow-water', 'swe', ''],
+              #['/home/6pz/data/PDEBench_Data/2D/NS_incom', 'incompNS', ''],
+              #['/home/6pz/data/PDEBench_Data/2D/CFD/2D_Train_Rand', compNS, '128'],
+              #['/home/6pz/data/PDEBench_Data/2D/CFD/2D_Train_Rand', compNS, '512'],
+              #['/home/6pz/data/PDEBench_Data/2D/CFD/2D_Train_Turb', compNS, ''],
+              #['/home/6pz/data/PDEBench_Data/2D/diffusion-reaction', 'diffre2d', ''],
+              ['/lustre/orion/lrn037/proj-shared/miniweather_thermals_tiny/test', 'thermalcollision2d', '', 'tk-2D'],
+              ]
+  append_datasets: [] # List of datasets to append to the input/output projections for finetuning
+

examples/config/Demo_MW_diffusion_avit.yaml

@@ -0,0 +1,88 @@
+basic_config: &basic_config
+  # Run settings
+  log_to_screen: !!bool True # Log progress to screen.
+  save_checkpoint: !!bool True # Save checkpoints
+  checkpoint_save_interval: 10 # Save every # epochs - also saves "best" according to val loss
+  true_time: !!bool False # Debugging setting - sets num workers to zero and activates syncs
+  num_data_workers: 1 # Generally pulling 8 cpu per process, so using 6 for DL - not sure if best ratio
+  enable_amp: !!bool False # Use automatic mixed precision - blows up with low variance fields right now
+  compile: !!bool False # Compile model - Does not currently work
+  gradient_checkpointing: !!bool False # Whether to use gradient checkpointing - Slow, but lower memory
+  exp_dir: 'Demo_Diffusion_MW_cond_avit_S_lt5' # Output path
+  # exp_dir: 'Demo_Diffusion_CIFAR10_correctedEpoch' # Output path
+  log_interval: 1 # How often to log - Don't think this is actually implemented
+  pretrained: !!bool False # Whether to load a pretrained model
+  # Training settings
+  drop_path: 0.1
+  batch_size: 32 #1
+  max_epochs: 200 #800
+  scheduler_epochs: -1
+  epoch_size: 100 #100 #2000 # Artificial epoch size
+  rescale_gradients: !!bool False # Activate hook that scales block gradients to norm 1
+  optimizer: 'AdamW' # a
+  scheduler: 'cosine' # Only cosine implemented
+  warmup_steps: 10 # Warmup when not using DAdapt
+  learning_rate: 1e-3 # -1 means use DAdapt
+  weight_decay: 1e-3
+  n_states: 15 #12   # Number of state variables across the datasets - Can be larger than real number and things will just go unused
+  state_names: ['Pressure', 'Vx', 'Vy', 'Density',  'Vx', 'Vy', 'Density', 'Pressure'] # Should be sorted
+  dt: 1 # Striding of data - Not currently implemented > 1
+  leadtime_max: 5 #prediction lead time range [1, leadtime_max]
+  n_steps: 1 # Length of history to include in input
+  enforce_max_steps: !!bool False # If false and n_steps > dataset steps, use dataset steps. Otherwise, raise Exception.
+  accum_grad: 1  # Real batch size is accum * batch_size, real steps/"epoch" is epoch_size / accum
+  # Model settings
+  # model_type: 'turbt'
+  # model_type:  'vit_all2all' # no need for time_type and space_type inputs
+  #model_type:  'svit' #currently only support time_type=="all2all_time" and space_type=="all2all"
+  #time_type: 'all2all_time' #
+  #space_type: 'all2all' #
+  model_type:  'avit' #currently only support space_type=="axial_attention" and time_type=="attention"
+  time_type: 'attention' # Conditional on block type
+  space_type: 'axial_attention' # Conditional on block type
+  tie_fields: !!bool False # Whether to use 1 embedding per field per data
+  embed_dim: 384 # Dimension of internal representation - 192/384/768/1024 for Ti/S/B/L
+  num_heads: 6 # Number of heads for attention - 3/6/12/16 for Ti/S/B/L
+  processor_blocks: 12 # Number of transformer blocks in the backbone - 12/12/12/24 for Ti/S/B/L
+  diffusion_config:
+    diffusion: !!bool True
+    cond_diffusion: !!bool True
+    model_channels: 128       # Base channel count for noise embedding MLP
+    channel_mult_emb: 4       # Multiplier for embedding MLP hidden dim (emb_channels = model_channels * channel_mult_emb)
+    embedding_type: 'positional'  # Noise level embedding type: 'positional' or 'fourier'
+    channel_mult_noise: 1     # Multiplier for noise embedding dim (noise_channels = model_channels * channel_mult_noise)
+  tokenizer_heads:
+    - head_name: "tk-2D"
+      patch_size: [[1, 4, 4]]
+  adaptive: !!bool False
+  sts_model: !!bool False
+  adap_samp: !!bool False
+  nrefines: 0 #number of coarse tokens picked to be refined
+  bias_type: 'none'  # Options rel, continuous, none
+  # Data settings
+  train_val_test: [.8, .1, .1]
+  augmentation: !!bool False # Augmentation not implemented
+  use_all_fields: !!bool True # Prepopulate the field metadata dictionary from dictionary in datasets
+  tie_batches: !!bool False # Force everything in batch to come from one dset
+  extended_names: !!bool False # Whether to use extended names - not currently implemented
+  embedding_offset: 0  # Use when adding extra finetuning fields
+  train_data_paths: [
+              #['/home/6pz/data/PDEBench_Data/2D/shallow-water', 'swe', ''],
+              #['/home/6pz/data/PDEBench_Data/2D/NS_incom', 'incompNS', ''],
+              #['/home/6pz/data/PDEBench_Data/2D/CFD/2D_Train_Rand', compNS, '128'],
+              #['/home/6pz/data/PDEBench_Data/2D/CFD/2D_Train_Rand', compNS, '512'],
+              #['/home/6pz/data/PDEBench_Data/2D/CFD/2D_Train_Turb', compNS, ''],
+              #['/home/6pz/data/PDEBench_Data/2D/diffusion-reaction', 'diffre2d', ''],
+              ['/lustre/orion/lrn037/proj-shared/miniweather_thermals_tiny/train', 'thermalcollision2d', '', 'tk-2D'],
+              ]
+  valid_data_paths: [
+              #['/home/6pz/data/PDEBench_Data/2D/shallow-water', 'swe', ''],
+              #['/home/6pz/data/PDEBench_Data/2D/NS_incom', 'incompNS', ''],
+              #['/home/6pz/data/PDEBench_Data/2D/CFD/2D_Train_Rand', compNS, '128'],
+              #['/home/6pz/data/PDEBench_Data/2D/CFD/2D_Train_Rand', compNS, '512'],
+              #['/home/6pz/data/PDEBench_Data/2D/CFD/2D_Train_Turb', compNS, ''],
+              #['/home/6pz/data/PDEBench_Data/2D/diffusion-reaction', 'diffre2d', ''],
+              ['/lustre/orion/lrn037/proj-shared/miniweather_thermals_tiny/test', 'thermalcollision2d', '', 'tk-2D'],
+              ]
+  append_datasets: [] # List of datasets to append to the input/output projections for finetuning
+

examples/config/Demo_MW_diffusion_svit.yaml

@@ -0,0 +1,88 @@
+basic_config: &basic_config
+  # Run settings
+  log_to_screen: !!bool True # Log progress to screen.
+  save_checkpoint: !!bool True # Save checkpoints
+  checkpoint_save_interval: 10 # Save every # epochs - also saves "best" according to val loss
+  true_time: !!bool False # Debugging setting - sets num workers to zero and activates syncs
+  num_data_workers: 1 # Generally pulling 8 cpu per process, so using 6 for DL - not sure if best ratio
+  enable_amp: !!bool False # Use automatic mixed precision - blows up with low variance fields right now
+  compile: !!bool False # Compile model - Does not currently work
+  gradient_checkpointing: !!bool False # Whether to use gradient checkpointing - Slow, but lower memory
+  exp_dir: 'Demo_Diffusion_MW_cond_svit_S_lt5' # Output path
+  # exp_dir: 'Demo_Diffusion_CIFAR10_correctedEpoch' # Output path
+  log_interval: 1 # How often to log - Don't think this is actually implemented
+  pretrained: !!bool False # Whether to load a pretrained model
+  # Training settings
+  drop_path: 0.1
+  batch_size: 32 #1
+  max_epochs: 200 #800
+  scheduler_epochs: -1
+  epoch_size: 100 #100 #2000 # Artificial epoch size
+  rescale_gradients: !!bool False # Activate hook that scales block gradients to norm 1
+  optimizer: 'AdamW' # a
+  scheduler: 'cosine' # Only cosine implemented
+  warmup_steps: 10 # Warmup when not using DAdapt
+  learning_rate: 1e-3 # -1 means use DAdapt
+  weight_decay: 1e-3
+  n_states: 15 #12   # Number of state variables across the datasets - Can be larger than real number and things will just go unused
+  state_names: ['Pressure', 'Vx', 'Vy', 'Density',  'Vx', 'Vy', 'Density', 'Pressure'] # Should be sorted
+  dt: 1 # Striding of data - Not currently implemented > 1
+  leadtime_max: 5 #prediction lead time range [1, leadtime_max]
+  n_steps: 1 # Length of history to include in input
+  enforce_max_steps: !!bool False # If false and n_steps > dataset steps, use dataset steps. Otherwise, raise Exception.
+  accum_grad: 1  # Real batch size is accum * batch_size, real steps/"epoch" is epoch_size / accum
+  # Model settings
+  # model_type: 'turbt'
+  # model_type:  'vit_all2all' # no need for time_type and space_type inputs
+  model_type:  'svit' #currently only support time_type=="all2all_time" and space_type=="all2all"
+  time_type: 'all2all_time' #
+  space_type: 'all2all' #
+  #model_type:  'avit' #currently only support space_type=="axial_attention" and time_type=="attention"
+  #time_type: 'attention' # Conditional on block type
+  #space_type: 'axial_attention' # Conditional on block type
+  tie_fields: !!bool False # Whether to use 1 embedding per field per data
+  embed_dim: 384 # Dimension of internal representation - 192/384/768/1024 for Ti/S/B/L
+  num_heads: 6 # Number of heads for attention - 3/6/12/16 for Ti/S/B/L
+  processor_blocks: 12 # Number of transformer blocks in the backbone - 12/12/12/24 for Ti/S/B/L
+  diffusion_config:
+    diffusion: !!bool True
+    cond_diffusion: !!bool True
+    model_channels: 128       # Base channel count for noise embedding MLP
+    channel_mult_emb: 4       # Multiplier for embedding MLP hidden dim (emb_channels = model_channels * channel_mult_emb)
+    embedding_type: 'positional'  # Noise level embedding type: 'positional' or 'fourier'
+    channel_mult_noise: 1     # Multiplier for noise embedding dim (noise_channels = model_channels * channel_mult_noise)
+  tokenizer_heads:
+    - head_name: "tk-2D"
+      patch_size: [[1, 4, 4]]
+  adaptive: !!bool False
+  sts_model: !!bool False
+  adap_samp: !!bool False
+  nrefines: 0 #number of coarse tokens picked to be refined
+  bias_type: 'none'  # Options rel, continuous, none
+  # Data settings
+  train_val_test: [.8, .1, .1]
+  augmentation: !!bool False # Augmentation not implemented
+  use_all_fields: !!bool True # Prepopulate the field metadata dictionary from dictionary in datasets
+  tie_batches: !!bool False # Force everything in batch to come from one dset
+  extended_names: !!bool False # Whether to use extended names - not currently implemented
+  embedding_offset: 0  # Use when adding extra finetuning fields
+  train_data_paths: [
+              #['/home/6pz/data/PDEBench_Data/2D/shallow-water', 'swe', ''],
+              #['/home/6pz/data/PDEBench_Data/2D/NS_incom', 'incompNS', ''],
+              #['/home/6pz/data/PDEBench_Data/2D/CFD/2D_Train_Rand', compNS, '128'],
+              #['/home/6pz/data/PDEBench_Data/2D/CFD/2D_Train_Rand', compNS, '512'],
+              #['/home/6pz/data/PDEBench_Data/2D/CFD/2D_Train_Turb', compNS, ''],
+              #['/home/6pz/data/PDEBench_Data/2D/diffusion-reaction', 'diffre2d', ''],
+              ['/lustre/orion/lrn037/proj-shared/miniweather_thermals_tiny/train', 'thermalcollision2d', '', 'tk-2D'],
+              ]
+  valid_data_paths: [
+              #['/home/6pz/data/PDEBench_Data/2D/shallow-water', 'swe', ''],
+              #['/home/6pz/data/PDEBench_Data/2D/NS_incom', 'incompNS', ''],
+              #['/home/6pz/data/PDEBench_Data/2D/CFD/2D_Train_Rand', compNS, '128'],
+              #['/home/6pz/data/PDEBench_Data/2D/CFD/2D_Train_Rand', compNS, '512'],
+              #['/home/6pz/data/PDEBench_Data/2D/CFD/2D_Train_Turb', compNS, ''],
+              #['/home/6pz/data/PDEBench_Data/2D/diffusion-reaction', 'diffre2d', ''],
+              ['/lustre/orion/lrn037/proj-shared/miniweather_thermals_tiny/test', 'thermalcollision2d', '', 'tk-2D'],
+              ]
+  append_datasets: [] # List of datasets to append to the input/output projections for finetuning
+