diff --git a/examples/basic_generate.py b/examples/basic_generate.py new file mode 100644 index 0000000..7b8c401 --- /dev/null +++ b/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) diff --git a/examples/config/Demo_MW_diffusion_TT.yaml b/examples/config/Demo_MW_diffusion_TT.yaml new file mode 100644 index 0000000..ec52378 --- /dev/null +++ b/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 + diff --git a/examples/config/Demo_MW_diffusion_avit.yaml b/examples/config/Demo_MW_diffusion_avit.yaml new file mode 100644 index 0000000..b50e150 --- /dev/null +++ b/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 + diff --git a/examples/config/Demo_MW_diffusion_svit.yaml b/examples/config/Demo_MW_diffusion_svit.yaml new file mode 100644 index 0000000..7c3a59c --- /dev/null +++ b/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 + diff --git a/examples/config/Demo_MW_diffusion_vit.yaml b/examples/config/Demo_MW_diffusion_vit.yaml new file mode 100644 index 0000000..736bf58 --- /dev/null +++ b/examples/config/Demo_MW_diffusion_vit.yaml @@ -0,0 +1,93 @@ +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_vit_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'], + # ['/lustre/orion/lrn037/proj-shared/CIFAR10/train/cifar10-32x32.zip', 'cifar10', 'tk-2D'], + # ['/lustre/orion/lrn037/proj-shared/CIFAR10_unzip/train', 'cifar10', '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'], + # ['/lustre/orion/lrn037/proj-shared/CIFAR10/test/cifar10-32x32.zip', 'cifar10', 'tk-2D'], + # ['/lustre/orion/lrn037/proj-shared/CIFAR10_unzip/test', 'cifar10', 'tk-2D'], + ] + append_datasets: [] # List of datasets to append to the input/output projections for finetuning + diff --git a/examples/submit_batch_diffusion.sh b/examples/submit_batch_diffusion.sh new file mode 100644 index 0000000..ad75d0e --- /dev/null +++ b/examples/submit_batch_diffusion.sh @@ -0,0 +1,43 @@ +#!/bin/bash +#SBATCH -A LRN037 +#SBATCH -J matey +#SBATCH -o %x-%j.out +#SBATCH -t 01:00:00 +#SBATCH -p batch +#SBATCH -N 1 +##SBATCH -q debug +#SBATCH -C nvme + +export OMP_NUM_THREADS=1 + +export master_node=$SLURMD_NODENAME +export config="basic_config" +export run_name="demo_diffusion" + +# export yaml_config=./config/Demo_MW_diffusion_avit.yaml +# export yaml_config=./config/Demo_MW_diffusion_svit.yaml +# export yaml_config=./config/Demo_MW_diffusion_vit.yaml +export yaml_config=./config/Demo_MW_diffusion_TT.yaml + + +##conda env with rocm 6.0.0 +#module load rocm/6.0.0 +#source /lustre/orion/proj-shared/lrn037/gounley1/conda600whl/etc/profile.d/conda.sh +#conda activate /lustre/orion/proj-shared/lrn037/gounley1/conda600whl +#module load cray-parallel-netcdf/1.12.3.9 + +##conda env with rocm 6.0.0 in world-shared +#source /lustre/orion/world-shared/lrn037/gounley1/env600.sh +source /lustre/orion/world-shared/stf218/junqi/forge/matey-env-rocm631.sh +export PYTHONPATH="${PYTHONPATH}:$(dirname "$PWD")" + +export MIOPEN_USER_DB_PATH=/mnt/bb/$USER/MIOPEN$SLURM_JOB_ID +export MIOPEN_CUSTOM_CACHE_DIR=${MIOPEN_USER_DB_PATH} +rm -rf ${MIOPEN_USER_DB_PATH} +mkdir -p ${MIOPEN_USER_DB_PATH} + +export MASTER_ADDR=$(hostname -i) +export MASTER_PORT=3442 + +srun -N$SLURM_JOB_NUM_NODES -n$((SLURM_JOB_NUM_NODES*8)) -c7 --gpu-bind=closest python basic_usage.py \ +--run_name $run_name --config $config --yaml_config $yaml_config --use_ddp diff --git a/examples/submit_batch_generate.sh b/examples/submit_batch_generate.sh new file mode 100644 index 0000000..4dbe598 --- /dev/null +++ b/examples/submit_batch_generate.sh @@ -0,0 +1,47 @@ +#!/bin/bash +#SBATCH -A LRN037 +#SBATCH -J matey +#SBATCH -o %x-%j.out +#SBATCH -t 00:15:00 +#SBATCH -p batch +#SBATCH -N 1 +##SBATCH -q debug +#SBATCH -C nvme + +export OMP_NUM_THREADS=1 + +export master_node=$SLURMD_NODENAME +export config="basic_config" + +# export model_dir="./Demo_Diffusion_MW_cond_avit_S_lt5/basic_config/demo_diffusion/" +# export model_dir="./Demo_Diffusion_MW_cond_svit_S_lt5/basic_config/demo_diffusion/" +# export model_dir="./Demo_Diffusion_MW_cond_vit_S_lt5/basic_config/demo_diffusion/" +export model_dir="./Demo_Diffusion_MW_cond_TT_3level_S_lt5/basic_config/demo_diffusion/" + +# export output_dir="./MW_cond_generation_outputs_batches_lt5/avit/" +# export output_dir="./MW_cond_generation_outputs_batches_lt5/svit/" +# export output_dir="./MW_cond_generation_outputs_batches_lt5/vit/" +export output_dir="./MW_cond_generation_outputs_batches_lt5/TT_3level/" + + +##conda env with rocm 6.0.0 +#module load rocm/6.0.0 +#source /lustre/orion/proj-shared/lrn037/gounley1/conda60s0whl/etc/profile.d/conda.sh +#conda activate /lustre/orion/proj-shared/lrn037/gounley1/conda600whl +#module load cray-parallel-netcdf/1.12.3.9 + +##conda env with rocm 6.0.0 in world-shared +#source /lustre/orion/world-shared/lrn037/gounley1/env600.sh +source /lustre/orion/world-shared/stf218/junqi/forge/matey-env-rocm631.sh +export PYTHONPATH="${PYTHONPATH}:$(dirname "$PWD")" + +export MIOPEN_USER_DB_PATH=/mnt/bb/$USER/MIOPEN$SLURM_JOB_ID +export MIOPEN_CUSTOM_CACHE_DIR=${MIOPEN_USER_DB_PATH} +rm -rf ${MIOPEN_USER_DB_PATH} +mkdir -p ${MIOPEN_USER_DB_PATH} + +export MASTER_ADDR=$(hostname -i) +export MASTER_PORT=3442 + +srun -N$SLURM_JOB_NUM_NODES -n$((SLURM_JOB_NUM_NODES*1)) -c7 --gpu-bind=closest python basic_generate.py \ +--config $config --model_dir $model_dir --output_dir $output_dir --use_ddp diff --git a/matey/__init__.py b/matey/__init__.py index 5bc0b97..f80468d 100644 --- a/matey/__init__.py +++ b/matey/__init__.py @@ -8,5 +8,5 @@ from .trustworthiness import * from .train import Trainer from .inference import Inferencer - +from .generate import Generator __version__ = "0.1.0" \ No newline at end of file diff --git a/matey/generate.py b/matey/generate.py new file mode 100644 index 0000000..59450f9 --- /dev/null +++ b/matey/generate.py @@ -0,0 +1,422 @@ +# from copyreg import pickle +import os +from random import seed +import torch +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP +from einops import rearrange +from collections import OrderedDict +from torchinfo import summary +from .data_utils.datasets import get_data_loader +from .models.avit import build_avit +from .models.svit import build_svit +from .models.vit import build_vit +from .models.turbt import build_turbt +from .models.diffusion_model import build_diffusion_model +from .utils.distributed_utils import determine_turt_levels, get_sequence_parallel_group +from .utils.forward_options import ForwardOptionsBase +from .utils.training_utils import EDMLoss +import json +import numpy as np +import pickle +import copy + + +class EDMSampler: + """ + EDM (Karras et al. 2022) 2nd-order Heun sampler. + + Encapsulates the denoising schedule and loop so that alternative samplers + can be dropped in by subclassing and overriding `sample`. + """ + def __init__(self, sigma_min=0.002, sigma_max=80, n_steps=18, rho=7, + S_churn=0, S_min=0, S_max=float('inf'), S_noise=1): + self.sigma_min = sigma_min + self.sigma_max = sigma_max + self.n_steps = n_steps + self.rho = rho + self.S_churn = S_churn + self.S_min = S_min + self.S_max = S_max + self.S_noise = S_noise + + def _t_steps(self, device): + idx = torch.arange(self.n_steps, device=device) + t = (self.sigma_max ** (1 / self.rho) + + idx / (self.n_steps - 1) + * (self.sigma_min ** (1 / self.rho) - self.sigma_max ** (1 / self.rho))) ** self.rho + return torch.cat([t, torch.zeros_like(t[:1])]) # t_N = 0 + + def _make_opts(self, opts_kwargs, blockdict, cond_dict_b): + return ForwardOptionsBase(**{**opts_kwargs, + 'blockdict': copy.deepcopy(blockdict), + 'cond_dict': copy.deepcopy(cond_dict_b)}) + + def sample(self, model, inp_b, num_samples, field_labels_b, bcs_b, + opts_kwargs, blockdict, cond_dict_b, + cond_diffusion=False, output_dir=None, batch_idx=None): + """ + Run the denoising loop over a batched input. + + inp_b : [T, B*num_samples, C, D, H, W] + Returns output grouped as [num_samples, T, B, C, D, H, W]. + """ + t_steps = self._t_steps(inp_b.device) + x_next = torch.randn_like(inp_b) * t_steps[0] + + for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])): + x_cur = x_next + gamma = (min(self.S_churn / self.n_steps, np.sqrt(2) - 1) + if self.S_min <= t_cur <= self.S_max else 0) + t_hat = t_cur + gamma * t_cur + x_hat = x_cur + (t_hat ** 2 - t_cur ** 2).sqrt() * self.S_noise * torch.randn_like(x_cur) + + # Euler step + opts = self._make_opts(opts_kwargs, blockdict, cond_dict_b) + if cond_diffusion: + opts.diffusion_cond = rearrange(inp_b, 't b c d h w -> b t c d h w') + denoised = model(x_hat, t_hat.repeat(x_hat.shape[1]), field_labels_b, bcs_b, opts) + d_cur = (x_hat - denoised) / t_hat + x_next = x_hat + (t_next - t_hat) * d_cur + + # 2nd-order Heun correction + if i < self.n_steps - 1: + opts = self._make_opts(opts_kwargs, blockdict, cond_dict_b) + if cond_diffusion: + opts.diffusion_cond = rearrange(inp_b, 't b c d h w -> b t c d h w') + denoised = model(x_next, t_next.repeat(x_next.shape[1]), field_labels_b, bcs_b, opts) + d_prime = (x_next - denoised) / t_next + x_next = x_hat + (t_next - t_hat) * (0.5 * d_cur + 0.5 * d_prime) + + if output_dir is not None: + grouped = rearrange(x_next, 't (b s) c d h w -> s t b c d h w', s=num_samples) + torch.save(grouped.cpu(), os.path.join(output_dir, f'generation_step_{i}_batch_{batch_idx}.pt')) + + return rearrange(x_next, 't (b s) c d h w -> s t b c d h w', s=num_samples) + + +class Generator: + def __init__(self, params, global_rank, local_rank, device): + self.device = device + self.params = params + self.global_rank = global_rank + self.local_rank = local_rank + self.world_size = int(os.environ.get("WORLD_SIZE", 1)) + self.log_to_screen = self.params.log_to_screen + # Basic setup + self.diffusion_loss = EDMLoss() + self.sampler = EDMSampler() + self.startEpoch = 0 + self.epoch = 0 + self.mp_type = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.half + + self.diffusion_config = getattr(self.params, 'diffusion_config', None) or {} + self.cond_diffusion = self.diffusion_config.get("cond_diffusion", False) + self.profiling = self.params.profiling if hasattr(self.params, "profiling") else False + + self.output_dir = self.params.output_dir + #define sequence parallel groups and local group info + if hasattr(self.params, "sp_groupsize"): + self.current_group, self.group_id, self.num_sequence_parallel_groups = get_sequence_parallel_group(sequence_parallel_groupsize=self.params.sp_groupsize) + else: + self.current_group, self.group_id, self.num_sequence_parallel_groups = get_sequence_parallel_group(num_sequence_parallel_groups=self.params.num_sequence_parallel_groups if hasattr(self.params, "num_sequence_parallel_groups") else self.world_size) + + self.group_rank = dist.get_rank(self.current_group) + self.group_size = dist.get_world_size(self.current_group) + + self.initialize_data() + #print(f"Initializing model on rank {self.global_rank}") + + #checking input_states value + labels_total=[self.train_dataset.subset_dict[dset] for dset in self.train_dataset.subset_dict] + labels_total = [item for sublist in labels_total for item in sublist] + if self.params.n_states t b c d h w') + except: + pass + + blockdict = getattr(self.valid_dataset.sub_dsets[dset_index[0]], "blockdict", None) + dset_type = self.valid_dataset.sub_dsets[dset_index[0]].type + tkhead_name = self.valid_dataset.sub_dsets[dset_index[0]].tkhead_name + imod = self.params.hierarchical["nlevels"]-1 if hasattr(self.params, "hierarchical") else 0 + if "graph" in data: + isgraph = True + inp = graphdata + imod_bottom = imod + else: + inp = rearrange(inp.to(self.device), 'b t c d h w -> t b c d h w') + isgraph = False + imod_bottom = determine_turt_levels(self.model.module.tokenizer_heads_params[tkhead_name][-1], inp.shape[-3:], imod) if imod>0 else 0 + + seq_group = self.current_group if dset_type in self.valid_dataset.DP_dsets else None + + torch.save(inp.cpu(), os.path.join(self.output_dir, f'inputdata_batch_{batch_idx}.pt')) + self.single_print(f'input saved to {os.path.join(self.output_dir, f"inputdata_batch_{batch_idx}.pt")}') + torch.save(tar.cpu(), os.path.join(self.output_dir, f'targetdata_batch_{batch_idx}.pt')) + self.single_print(f'target saved to {os.path.join(self.output_dir, f"targetdata_batch_{batch_idx}.pt")}') + + torch.save(torch.tensor(leadtime).cpu(), os.path.join(self.output_dir, f'leadtimedata_batch_{batch_idx}.pt')) + self.single_print(f'leadtime saved to {os.path.join(self.output_dir, f"leadtimedata_batch_{batch_idx}.pt")}') + + self.model.eval() + + with torch.no_grad(): + # Expand batch dimension to run all num_samples in one forward pass. + # inp: [T, B, C, D, H, W] -> [T, B*num_samples, C, D, H, W] + inp_b = inp.repeat_interleave(num_samples, dim=1) + field_labels_b = field_labels.repeat_interleave(num_samples, dim=0) + bcs_b = bcs.repeat_interleave(num_samples, dim=0) + leadtime_b = leadtime.repeat_interleave(num_samples, dim=0) if leadtime is not None else None + cond_input_b = cond_input.repeat_interleave(num_samples, dim=0) if cond_input is not None else None + cond_dict_b = {} + if cond_dict: + cond_dict_b["labels"] = cond_dict["labels"].repeat_interleave(num_samples, dim=0) + cond_dict_b["fields"] = cond_dict["fields"].repeat_interleave(num_samples, dim=1) + + opts_kwargs = dict( + imod=imod, + imod_bottom=imod_bottom, + tkhead_name=tkhead_name, + sequence_parallel_group=seq_group, + leadtime=leadtime_b, + cond_input=cond_input_b, + isgraph=isgraph, + field_labels_out=field_labels_b, + ) + output = self.sampler.sample( + self.model, inp_b, num_samples, field_labels_b, bcs_b, + opts_kwargs=opts_kwargs, + blockdict=blockdict, + cond_dict_b=cond_dict_b, + cond_diffusion=self.cond_diffusion, + output_dir=self.output_dir, + batch_idx=batch_idx, + ) + torch.save(output.cpu(), os.path.join(self.output_dir, f'generation_output_batch_{batch_idx}.pt')) + self.single_print(f'Generation output saved to {os.path.join(self.output_dir, f"generation_output_batch_{batch_idx}.pt")}') + + torch.cuda.empty_cache() + + + + def autoregressive_generate(self, seed=None, num_samples=1, num_steps=10): + ### FIXME: UNTESTED ### + if self.global_rank == 0: + summary(self.model) + self.single_print("Starting Autoregressive Generation Loop...") + + if seed is not None: + torch.manual_seed(seed) + + # Load one batch of data + data_iter = iter(self.valid_data_loader) + data = next(data_iter) + + if "graph" in data: + graphdata = data["graph"].to(self.device) + tar = graphdata.y + leadtime = graphdata.leadtime + dset_index, field_labels, field_labels_out, bcs = map(lambda x: x.to(self.device), [data[varname] for varname in ["dset_idx", "field_labels", "field_labels_out", "bcs"]]) + else: + inp, dset_index, field_labels, bcs, tar, leadtime = map(lambda x: x.to(self.device), [data[varname] for varname in ["input", "dset_idx", "field_labels", "bcs", "label", "leadtime"]]) + field_labels_out = field_labels + supportdata = True if hasattr(self.params, 'supportdata') else False + cond_input = data["cond_input"].to(self.device) if supportdata else None + + cond_dict = {} + try: + cond_dict["labels"] = data["cond_field_labels"].to(self.device) + cond_dict["fields"] = rearrange(data["cond_fields"].to(self.device), 'b t c d h w -> t b c d h w') + except: + pass + + blockdict = getattr(self.valid_dataset.sub_dsets[dset_index[0]], "blockdict", None) + dset_type = self.valid_dataset.sub_dsets[dset_index[0]].type + tkhead_name = self.valid_dataset.sub_dsets[dset_index[0]].tkhead_name + imod = self.params.hierarchical["nlevels"]-1 if hasattr(self.params, "hierarchical") else 0 + if "graph" in data: + isgraph = True + inp = graphdata + imod_bottom = imod + else: + inp = rearrange(inp.to(self.device), 'b t c d h w -> t b c d h w') + isgraph = False + imod_bottom = determine_turt_levels(self.model.module.tokenizer_heads_params[tkhead_name][-1], inp.shape[-3:], imod) if imod > 0 else 0 + + seq_group = self.current_group if dset_type in self.valid_dataset.DP_dsets else None + + # Save initial conditioning input and target once + torch.save(inp.cpu(), os.path.join(self.output_dir, 'inputdata_step0.pt')) + self.single_print(f'input saved to {os.path.join(self.output_dir, "inputdata_step0.pt")}') + torch.save(tar.cpu(), os.path.join(self.output_dir, 'targetdata_step0.pt')) + self.single_print(f'target saved to {os.path.join(self.output_dir, "targetdata_step0.pt")}') + torch.save(torch.tensor(leadtime).cpu(), os.path.join(self.output_dir, 'leadtimedata_step0.pt')) + self.single_print(f'leadtime saved to {os.path.join(self.output_dir, "leadtimedata_step0.pt")}') + + self.model.eval() + inp_cur = inp.clone() # conditioning that advances each autoregressive step + + for step_idx in range(num_steps): + with torch.no_grad(): + # Expand batch dimension to run all num_samples in one forward pass. + # inp_cur: [T, B, C, D, H, W] -> [T, B*num_samples, C, D, H, W] + inp_b = inp_cur.repeat_interleave(num_samples, dim=1) + field_labels_b = field_labels.repeat_interleave(num_samples, dim=0) + bcs_b = bcs.repeat_interleave(num_samples, dim=0) + leadtime_b = leadtime.repeat_interleave(num_samples, dim=0) if leadtime is not None else None + cond_input_b = cond_input.repeat_interleave(num_samples, dim=0) if cond_input is not None else None + cond_dict_b = {} + if cond_dict: + cond_dict_b["labels"] = cond_dict["labels"].repeat_interleave(num_samples, dim=0) + cond_dict_b["fields"] = cond_dict["fields"].repeat_interleave(num_samples, dim=1) + + opts_kwargs = dict( + imod=imod, + imod_bottom=imod_bottom, + tkhead_name=tkhead_name, + sequence_parallel_group=seq_group, + leadtime=leadtime_b, + cond_input=cond_input_b, + isgraph=isgraph, + field_labels_out=field_labels_b, + ) + # output: [num_samples, T, B, C, D, H, W] + output = self.sampler.sample( + self.model, inp_b, num_samples, field_labels_b, bcs_b, + opts_kwargs=opts_kwargs, + blockdict=blockdict, + cond_dict_b=cond_dict_b, + cond_diffusion=self.cond_diffusion, + ) + + torch.save(output.cpu(), os.path.join(self.output_dir, f'generation_output_step_{step_idx}.pt')) + self.single_print(f'Generation output saved to {os.path.join(self.output_dir, f"generation_output_step_{step_idx}.pt")}') + + # Mean across samples becomes the conditioning for the next autoregressive step + inp_cur = output.mean(dim=0) # [T, B, C, D, H, W] + torch.save(inp_cur.cpu(), os.path.join(self.output_dir, f'generation_mean_step_{step_idx}.pt')) + self.single_print(f'Step {step_idx} ensemble mean saved to {os.path.join(self.output_dir, f"generation_mean_step_{step_idx}.pt")}') + + torch.cuda.empty_cache() + + diff --git a/matey/models/__init__.py b/matey/models/__init__.py index fae8905..92197af 100644 --- a/matey/models/__init__.py +++ b/matey/models/__init__.py @@ -6,5 +6,4 @@ from .svit import build_svit, sViT_all2all from .vit import build_vit, ViT_all2all from .turbt import build_turbt, TurbT - __all__ = ["build_avit", "build_svit", "build_vit","build_turbt", "AViT","sViT_all2all","ViT_all2all","TurbT"] diff --git a/matey/models/attention_modules.py b/matey/models/attention_modules.py index fe825c8..8757141 100644 --- a/matey/models/attention_modules.py +++ b/matey/models/attention_modules.py @@ -307,6 +307,10 @@ def forward(self, x, sequence_parallel_group=None, leadtime=None, local_att=Fals #leadtime: btoken_len x c #t_pos_area: token_len x 4 B, C, len = x.shape + + # Bypass attention if sequence length is 1 + if len == 1: + return x input = x.clone() # Rearrange and prenorm diff --git a/matey/models/avit.py b/matey/models/avit.py index 9ff2eea..3fe909f 100644 --- a/matey/models/avit.py +++ b/matey/models/avit.py @@ -36,7 +36,8 @@ def build_avit(params): cond_input=getattr(params,'supportdata', False), n_steps=params.n_steps, bias_type=params.bias_type, - hierarchical=getattr(params, 'hierarchical', None) + hierarchical=getattr(params, 'hierarchical', None), + diffusion_config=getattr(params, 'diffusion_config', None), ) return model @@ -52,9 +53,11 @@ class AViT(BaseModel): n_states (int): Number of input state variables. """ def __init__(self, tokenizer_heads=None, embed_dim=768, space_type="axial_attention", time_type="attention", num_heads=12, processor_blocks=8, n_states=6, n_states_cond=None, - drop_path=.2, sts_train=False, sts_model=False, leadtime=False, cond_input=False, n_steps=1, bias_type="none", SR_ratio=[1,1,1], hierarchical=None): + drop_path=.2, sts_train=False, sts_model=False, leadtime=False, cond_input=False, n_steps=1, bias_type="none", SR_ratio=[1,1,1], hierarchical=None, + diffusion_config=None): super().__init__(tokenizer_heads=tokenizer_heads, n_states=n_states, n_states_cond=n_states_cond, embed_dim=embed_dim, leadtime=leadtime, - cond_input=cond_input, n_steps=n_steps, bias_type=bias_type, SR_ratio=SR_ratio, hierarchical=hierarchical) + cond_input=cond_input, n_steps=n_steps, bias_type=bias_type, SR_ratio=SR_ratio, hierarchical=hierarchical, + diffusion_config=diffusion_config) self.drop_path = drop_path self.dp = np.linspace(0, drop_path, processor_blocks) @@ -165,16 +168,25 @@ def forward(self, x, state_labels, bcs, opts: ForwardOptionsBase, train_opts: Op isgraph = opts.isgraph ################################################################## conditioning = (cond_dict != None and bool(cond_dict) and self.conditioning) + sigma = getattr(opts, 'sigma', None) + diffusion_cond = getattr(opts, 'diffusion_cond', None) + persample_normalize = not self.diffusion + ################################################################## assert not isgraph, "graph is not supported in AViT" #T,B,C,D,H,W T, _, _, D, _, _ = x.shape + + if self.diffusion: + emb = self.compute_diffusion_emb(sigma, imod) + if diffusion_cond is not None: + diffusion_cond = rearrange(diffusion_cond, 'b t c d h w -> t b c d h w') if self.tokenizer_heads_gammaref[tkhead_name] is None and refine_ratio is None: refineind=None else: refineind = get_top_variance_patchids(self.tokenizer_heads_params[tkhead_name], x, self.tokenizer_heads_gammaref[tkhead_name], refine_ratio) - x, data_mean, data_std = normalize_spatiotemporal_persample(x) - #self.debug_nan(x, message="input") + if persample_normalize: + x, data_mean, data_std = normalize_spatiotemporal_persample(x) #[T, B, C_emb//4, D, H, W] x_pre = self.get_unified_preembedding(x, state_labels, self.space_bag[imod]) @@ -202,6 +214,21 @@ def forward(self, x, state_labels, bcs, opts: ForwardOptionsBase, train_opts: Op c_pre = self.get_unified_preembedding(cond_dict["fields"], cond_dict["labels"], self.space_bag_cond[imod]) c = self.get_structured_sequence(c_pre, -1, self.tokenizer_ensemble_heads[imod][tkhead_name]["embed_cond"]) + if diffusion_cond is not None: + diffusion_cond_pre = self.get_unified_preembedding( + diffusion_cond, state_labels, self.space_bag[imod]) + diffusion_cond_tokens = self.get_structured_sequence( + diffusion_cond_pre, -1, self.tokenizer_ensemble_heads[imod][tkhead_name]["embed"]) + # diffusion_cond_tokens: t b embed_dim ntokD ntokH ntokW + _, B_x, _, D_x, H_x, W_x = x.shape + cond_combined = diffusion_cond_tokens + rearrange(emb, 'b c -> 1 b c 1 1 1') + cond_fused = self.diffusion_cond_proj[str(imod)]( + rearrange(cond_combined, 't b c d h w -> (t b d h w) c')) + x = x + rearrange(cond_fused, '(t b d h w) c -> t b c d h w', + t=T, b=B_x, d=D_x, h=H_x, w=W_x) + elif self.diffusion: + x = x + rearrange(emb, 'b c -> 1 b c 1 1 1') + # Process for iblk, blk in enumerate(self.blocks): if conditioning: @@ -228,8 +255,10 @@ def forward(self, x, state_labels, bcs, opts: ForwardOptionsBase, train_opts: Op x = self.add_sts_model(x, x_pre, state_labels, bcs, tkhead_name, leadtime=leadtime, refineind=refineind, blockdict=blockdict) # Denormalize - x = x * data_std + data_mean # All state labels in the batch should be identical + if persample_normalize: + x = x * data_std + data_mean if train_opts is not None and train_opts.returnbase4train: - xbase = xbase * data_std + data_mean + if persample_normalize: + xbase = xbase * data_std + data_mean return x[-1], xbase[-1] return x[-1] # Just return last step - now just predict delta. diff --git a/matey/models/basemodel.py b/matey/models/basemodel.py index 01d9ff9..3b7d3d7 100644 --- a/matey/models/basemodel.py +++ b/matey/models/basemodel.py @@ -8,7 +8,7 @@ import numpy as np from einops import rearrange, repeat from .spatial_modules import hMLP_stem, hMLP_output, SubsampledLinear, GraphhMLP_stem, GraphhMLP_output, UpsampleinSpace, UpsampleConv3d -from .time_modules import leadtimeMLP +from .time_modules import leadtimeMLP, FourierEmbedding_EDM, PositionalEmbedding_EDM, Linear_EDM from .input_modules import input_control_MLP from .positionbias_modules import positionbias_mod import sys @@ -29,7 +29,7 @@ class BaseModel(nn.Module): n_states (int): Number of input state variables. """ def __init__(self, tokenizer_heads, n_states=6, n_states_out=None, n_states_cond=None, embed_dim=768, leadtime=False, cond_input=False, n_steps=1, bias_type="none", SR_ratio=[1,1,1], model_SR=False, - hierarchical=None, notransposed=False, nlevels=1, smooth=False, use_linear=False, ghost_sync=False): + hierarchical=None, notransposed=False, nlevels=1, smooth=False, use_linear=False, ghost_sync=False, diffusion_config=None): super().__init__() self.space_bag = nn.ModuleList([SubsampledLinear(n_states, embed_dim//4) for _ in range(nlevels)]) self.conditioning = (n_states_cond is not None and n_states_cond > 0) @@ -91,7 +91,32 @@ def __init__(self, tokenizer_heads, n_states=6, n_states_out=None, n_states_cond if self.cond_input: self.inconMLP.append(input_control_MLP(hidden_dim=embed_dim,n_steps=n_steps)) self.posbias.append(positionbias_mod(bias_type, embed_dim)) - self.embed_dim=embed_dim + _dc = diffusion_config if diffusion_config is not None else {} + self.diffusion = _dc.get('diffusion', False) + if self.diffusion: + model_channels = _dc.get('model_channels', 128) + channel_mult_emb = _dc.get('channel_mult_emb', 4) + embedding_type = _dc.get('embedding_type', 'positional') + channel_mult_noise = _dc.get('channel_mult_noise', 1) + emb_channels = model_channels * channel_mult_emb + noise_channels = model_channels * channel_mult_noise + init = dict(init_mode='xavier_uniform') + + self.map_noise = (PositionalEmbedding_EDM(num_channels=noise_channels, endpoint=True) + if embedding_type == 'positional' + else FourierEmbedding_EDM(num_channels=noise_channels)) + self.map_layer0 = Linear_EDM(in_features=noise_channels, out_features=emb_channels, **init) + self.map_layer1 = Linear_EDM(in_features=emb_channels, out_features=emb_channels, **init) + + self.affine = nn.ModuleDict({}) + for ilevel in range(nlevels): + self.affine[str(ilevel)] = Linear_EDM(in_features=emb_channels, out_features=embed_dim, **init) + + self.diffusion_cond_proj = nn.ModuleDict({}) + for ilevel in range(nlevels): + self.diffusion_cond_proj[str(ilevel)] = Linear_EDM(in_features=embed_dim, out_features=embed_dim, **init) + + self.embed_dim=embed_dim def expand_conv_projections(self, refine_resol): """ Appends addition conv heads""" @@ -680,9 +705,16 @@ def get_spatiotemporalfromsequence(self, x_padding, patch_ids, patch_ids_ref, sp x_reconst = rearrange(x_coarsen, '(b ntz ntx nty) t c d h w -> t b c (ntz d) (ntx h) (nty w)', ntz=ntokendim[0], ntx=ntokendim[1], nty=ntokendim[2]) return x_reconst + def compute_diffusion_emb(self, sigma, imod): + emb = self.map_noise(sigma) + emb = emb.reshape(emb.shape[0], 2, -1).flip(1).reshape(*emb.shape) + emb = F.silu(self.map_layer0(emb)) + emb = F.silu(self.map_layer1(emb)) + return self.affine[str(imod)](emb) + def add_sts_model(self): pass - + def forward(self): raise NotImplementedError diff --git a/matey/models/diffusion_model.py b/matey/models/diffusion_model.py new file mode 100644 index 0000000..5c67b71 --- /dev/null +++ b/matey/models/diffusion_model.py @@ -0,0 +1,73 @@ +import torch +import torch.nn as nn + +from .avit import build_avit +from .svit import build_svit +from .vit import build_vit +from .turbt import build_turbt + +def build_diffusion_model(params): + model = EDMPrecond(params=params) + return model + + +class EDMPrecond(nn.Module): + def __init__(self, + label_dim = 0, # Number of class labels, 0 = unconditional. + use_fp16 = False, # Execute the underlying model at FP16 precision? + sigma_min = 0, # Minimum supported noise level. + sigma_max = float('inf'), # Maximum supported noise level. + sigma_data = 0.5, # Expected standard deviation of the training data. + params = None, # Additional parameters for the underlying model. + ): + super().__init__() + + self.label_dim = label_dim + self.use_fp16 = use_fp16 + self.sigma_min = sigma_min + self.sigma_max = sigma_max + self.sigma_data = sigma_data + + if params.model_type == 'avit': + self.model = build_avit(params) + self.tokenizer_heads_params = self.model.tokenizer_heads_params + elif params.model_type == 'svit': + self.model = build_svit(params) + self.tokenizer_heads_params = self.model.tokenizer_heads_params + elif params.model_type == 'vit_all2all': + self.model = build_vit(params) + self.tokenizer_heads_params = self.model.tokenizer_heads_params + elif params.model_type == 'turbt': + self.model = build_turbt(params) + self.tokenizer_heads_params = self.model.tokenizer_heads_params + else: + raise ValueError(f"Unknown diffusion model type: {params.model_type}") + + + def forward(self, x, sigma, field_labels, bcs, opts, force_fp32=False, **model_kwargs): + x = x.to(torch.float32) + sigma = sigma.to(torch.float32).reshape(1, -1, 1, 1, 1, 1) # reshape for broadcasting. shape: [1, B, 1, 1, 1, 1] + # class_labels = None if self.label_dim == 0 else torch.zeros([1, self.label_dim], device=x.device) if class_labels is None else class_labels.to(torch.float32).reshape(-1, self.label_dim) + dtype = torch.float16 if (self.use_fp16 and not force_fp32 and x.device.type == 'cuda') else torch.float32 + + c_skip = self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2) + c_out = sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2).sqrt() + c_in = 1 / (self.sigma_data ** 2 + sigma ** 2).sqrt() + c_noise = sigma.log() / 4 + + opts.sigma = c_noise.flatten() + # print(f"opts.imod: {opts.imod}, opts.imod_bottom: {opts.imod_bottom}") + # print(f"x shape: {x.shape}, c_in shape: {c_in.shape}, c_skip shape: {c_skip.shape}, c_out shape: {c_out.shape}, c_noise shape: {opts.sigma.shape}") + # print(f"class_labels shape: {opts.diffusion_cond.shape}") + F_x = self.model((c_in * x).to(dtype), field_labels, bcs, opts) + + F_x = F_x.unsqueeze(0) + # print(f"x shape: {x.shape}, F_x shape: {F_x.shape}, c_in shape: {c_in.shape}, c_skip shape: {c_skip.shape}, c_out shape: {c_out.shape}, c_noise shape: {opts.sigma.shape}") + + assert F_x.dtype == dtype + D_x = c_skip * x + c_out * F_x.to(torch.float32) + + # print(f"x shape: {x.shape}, c_skip shape: {c_skip.shape}, c_out shape: {c_out.shape}, F_x shape: {F_x.shape}, D_x shape: {D_x.shape}") + + return D_x + diff --git a/matey/models/svit.py b/matey/models/svit.py index c1fdfe4..3cdd1a5 100644 --- a/matey/models/svit.py +++ b/matey/models/svit.py @@ -6,6 +6,7 @@ import torch.nn as nn import numpy as np from einops import rearrange + from .spacetime_modules import SpaceTimeBlock_svit from .basemodel import BaseModel from ..data_utils.shared_utils import normalize_spatiotemporal_persample, get_top_variance_patchids, normalize_spatiotemporal_persample_graph @@ -39,7 +40,8 @@ def build_svit(params): bias_type=params.bias_type, replace_patch=getattr(params, 'replace_patch', True), ghost_sync=getattr(params, 'ghost_sync', False), - hierarchical=getattr(params, 'hierarchical', None) + hierarchical=getattr(params, 'hierarchical', None), + diffusion_config=getattr(params, 'diffusion_config', None), ) return model @@ -54,9 +56,11 @@ class sViT_all2all(BaseModel): sts_f """ def __init__(self, tokenizer_heads=None, embed_dim=768, space_type="all2all", time_type="all2all", num_heads=12, processor_blocks=8, n_states=6, n_states_cond=None, ghost_sync=False, - drop_path=.2, sts_train=False, sts_model=False, leadtime=False, cond_input=False, n_steps=1, bias_type="none", replace_patch=True, SR_ratio=[1,1,1], hierarchical=None): - super().__init__(tokenizer_heads=tokenizer_heads, n_states=n_states, n_states_cond=n_states_cond, embed_dim=embed_dim, leadtime=leadtime, - cond_input=cond_input, n_steps=n_steps, bias_type=bias_type, SR_ratio=SR_ratio, hierarchical=hierarchical, ghost_sync=ghost_sync) + drop_path=.2, sts_train=False, sts_model=False, leadtime=False, cond_input=False, n_steps=1, bias_type="none", replace_patch=True, SR_ratio=[1,1,1], hierarchical=None, + diffusion_config=None): + super().__init__(tokenizer_heads=tokenizer_heads, n_states=n_states, n_states_cond=n_states_cond, embed_dim=embed_dim, leadtime=leadtime, + cond_input=cond_input, n_steps=n_steps, bias_type=bias_type, SR_ratio=SR_ratio, hierarchical=hierarchical, ghost_sync=ghost_sync, + diffusion_config=diffusion_config) self.drop_path = drop_path self.dp = np.linspace(0, drop_path, processor_blocks) @@ -136,12 +140,20 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase, train_opts: isgraph=opts.isgraph field_labels_out=opts.field_labels_out ghost_info = opts.ghost_info + sigma = getattr(opts, 'sigma', None) + diffusion_cond = getattr(opts, 'diffusion_cond', None) + persample_normalize = not self.diffusion ################################################################## conditioning = (cond_dict != None and bool(cond_dict) and self.conditioning) if field_labels_out is None: field_labels_out = state_labels + if self.diffusion: + emb = self.compute_diffusion_emb(sigma, imod) + if diffusion_cond is not None: + diffusion_cond = rearrange(diffusion_cond, 'b t c d h w -> t b c d h w') + if isgraph: with torch.no_grad(): check_same_sample_across_halo(data, ghost_info, sequence_parallel_group) if sequence_parallel_group is not None else None @@ -159,8 +171,8 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase, train_opts: refineind=None else: refineind = get_top_variance_patchids(self.tokenizer_heads_params[tkhead_name], x, self.tokenizer_heads_gammaref[tkhead_name], refine_ratio) - #self.debug_nan(x, message="input") - x, data_mean, data_std = normalize_spatiotemporal_persample(x) + if persample_normalize: + x, data_mean, data_std = normalize_spatiotemporal_persample(x) ################################################################################ if self.leadtime and leadtime is not None: leadtime = self.ltimeMLP[imod](leadtime) @@ -184,6 +196,19 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase, train_opts: assert refineind == None assert not isgraph, "Not set conditioning yet" c, _, _, _, _, _, _, _ = self.get_patchsequence(cond_dict["fields"], cond_dict["labels"], tkhead_name, refineind=refineind, blockdict=blockdict, ilevel=imod, conditioning=conditioning) + + if diffusion_cond is not None: + diffusion_cond_padding, _, _, _, _, _, _, _ = \ + self.get_patchsequence(diffusion_cond, state_labels, tkhead_name, + refineind=None, blockdict=blockdict, ilevel=imod, isgraph=isgraph) + # diffusion_cond_padding: t b c ntoken_tot + cond_combined = diffusion_cond_padding + rearrange(emb, 'b c -> 1 b c 1') + cond_fused = self.diffusion_cond_proj[str(imod)]( + rearrange(cond_combined, 't b c L -> (t b L) c')) + x_padding = x_padding + rearrange(cond_fused, '(t b L) c -> t b c L', t=T, b=x_padding.shape[1]) + elif self.diffusion: + x_padding = x_padding + rearrange(emb, 'b c -> 1 b c 1') + ################################################################################ if self.posbias[imod] is not None and tposarea_padding is not None: posbias = self.posbias[imod](tposarea_padding, mask_padding=mask_padding, use_zpos=True if D>1 else False) # b t L c->b t L c_emb @@ -231,10 +256,12 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase, train_opts: return x ######### Denormalize ######## x = x[:,:,field_labels_out[0],...] - x = x * data_std + data_mean + if persample_normalize: + x = x * data_std + data_mean ################################################################################ if train_opts is not None and train_opts.returnbase4train: - xbase = xbase * data_std + data_mean + if persample_normalize: + xbase = xbase * data_std + data_mean return x[-1], xbase[-1] return x[-1] diff --git a/matey/models/time_modules.py b/matey/models/time_modules.py index c903c81..011769b 100644 --- a/matey/models/time_modules.py +++ b/matey/models/time_modules.py @@ -3,6 +3,8 @@ # This file is part of the MATEY Project. import torch.nn as nn +import torch +import numpy as np class leadtimeMLP(nn.Module): def __init__(self, hidden_dim, exp_factor=4.): @@ -14,3 +16,52 @@ def __init__(self, hidden_dim, exp_factor=4.): def forward(self, x): return self.fc2(self.act(self.fc1(x))) + +#From https://github.com/NVlabs/edm/blob/008a4e5316c8e3bfe61a62f874bddba254295afb/training/networks.py#L193 +#copied for now, need to figure out if edm can be installed as a package and if we can use these functions like edm.XX [Nov, 2025] +class PositionalEmbedding_EDM(torch.nn.Module): + def __init__(self, num_channels, max_positions=10000, endpoint=False): + super().__init__() + self.num_channels = num_channels + self.max_positions = max_positions + self.endpoint = endpoint + + def forward(self, x): + freqs = torch.arange(start=0, end=self.num_channels//2, dtype=torch.float32, device=x.device) + freqs = freqs / (self.num_channels // 2 - (1 if self.endpoint else 0)) + freqs = (1 / self.max_positions) ** freqs + x = x.ger(freqs.to(x.dtype)) + x = torch.cat([x.cos(), x.sin()], dim=1) + return x + +class FourierEmbedding_EDM(torch.nn.Module): + def __init__(self, num_channels, scale=16): + super().__init__() + self.register_buffer('freqs', torch.randn(num_channels // 2) * scale) + + def forward(self, x): + x = x.ger((2 * np.pi * self.freqs).to(x.dtype)) + x = torch.cat([x.cos(), x.sin()], dim=1) + return x + +def weight_init(shape, mode, fan_in, fan_out): + if mode == 'xavier_uniform': return np.sqrt(6 / (fan_in + fan_out)) * (torch.rand(*shape) * 2 - 1) + if mode == 'xavier_normal': return np.sqrt(2 / (fan_in + fan_out)) * torch.randn(*shape) + if mode == 'kaiming_uniform': return np.sqrt(3 / fan_in) * (torch.rand(*shape) * 2 - 1) + if mode == 'kaiming_normal': return np.sqrt(1 / fan_in) * torch.randn(*shape) + raise ValueError(f'Invalid init mode "{mode}"') + +class Linear_EDM(torch.nn.Module): + def __init__(self, in_features, out_features, bias=True, init_mode='kaiming_normal', init_weight=1, init_bias=0): + super().__init__() + self.in_features = in_features + self.out_features = out_features + init_kwargs = dict(mode=init_mode, fan_in=in_features, fan_out=out_features) + self.weight = torch.nn.Parameter(weight_init([out_features, in_features], **init_kwargs) * init_weight) + self.bias = torch.nn.Parameter(weight_init([out_features], **init_kwargs) * init_bias) if bias else None + + def forward(self, x): + x = x @ self.weight.to(x.dtype).t() + if self.bias is not None: + x = x.add_(self.bias.to(x.dtype)) + return x diff --git a/matey/models/turbt.py b/matey/models/turbt.py index 7798d23..4fdbe4d 100644 --- a/matey/models/turbt.py +++ b/matey/models/turbt.py @@ -12,6 +12,7 @@ from ..data_utils.shared_utils import normalize_spatiotemporal_persample, get_top_variance_patchids, normalize_spatiotemporal_persample_graph from ..data_utils.utils import construct_filterkernel, construct_filterkernel2D from .spatial_modules import UpsampleinSpace + import sys, copy from operator import mul from functools import reduce @@ -43,7 +44,8 @@ def build_turbt(params): replace_patch=getattr(params, 'replace_patch', True), hierarchical=getattr(params, 'hierarchical', None), ghost_sync=getattr(params, 'ghost_sync', False), - notransposed=getattr(params, 'notransposed', False) + notransposed=getattr(params, 'notransposed', False), + diffusion_config=getattr(params, 'diffusion_config', None), ) return model @@ -60,9 +62,11 @@ class TurbT(BaseModel): sts_f """ def __init__(self, tokenizer_heads=None, embed_dim=768, num_heads=12, processor_blocks=8, n_states=6, ghost_sync=False, - drop_path=.2, sts_train=False, sts_model=False, leadtime=False, cond_input=False, n_steps=1, bias_type="none", replace_patch=True, hierarchical=None, notransposed=False): - super().__init__(tokenizer_heads=tokenizer_heads, n_states=n_states, embed_dim=embed_dim, leadtime=leadtime, cond_input=cond_input, n_steps=n_steps, bias_type=bias_type, hierarchical=hierarchical, - notransposed=notransposed, nlevels=hierarchical["nlevels"] if hierarchical is not None else 1, ghost_sync=ghost_sync) + drop_path=.2, sts_train=False, sts_model=False, leadtime=False, cond_input=False, n_steps=1, bias_type="none", replace_patch=True, hierarchical=None, notransposed=False, + diffusion_config=None): + super().__init__(tokenizer_heads=tokenizer_heads, n_states=n_states, embed_dim=embed_dim, leadtime=leadtime, cond_input=cond_input, n_steps=n_steps, bias_type=bias_type, hierarchical=hierarchical, + notransposed=notransposed, nlevels=hierarchical["nlevels"] if hierarchical is not None else 1, ghost_sync=ghost_sync, + diffusion_config=diffusion_config) self.drop_path = drop_path self.dp = np.linspace(0, drop_path, processor_blocks) self.module_blocks=nn.ModuleDict({}) @@ -118,7 +122,7 @@ def __init__(self, tokenizer_heads=None, embed_dim=768, num_heads=12, processor #FIXME: figure out how to do local attention in 3D physical space for corrections self.module_blocks[str(imod)] = nn.ModuleList([SpaceTimeBlock_all2all(embed_dim, num_heads,drop_path=self.dp[i]) for i in range(processor_blocks//self.nhlevels)]) - + def filterdata(self, data, blockdict=None): #T,B,C,D,H,W assert data.ndim==6, f"unkown tensor shape in filter_data, {data.shape}" @@ -225,15 +229,25 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase): isgraph=opts.isgraph field_labels_out=opts.field_labels_out ghost_info = opts.ghost_info + sigma = getattr(opts, 'sigma', None) + diffusion_cond = getattr(opts, 'diffusion_cond', None) + persample_normalize = not self.diffusion ################################################################## if refine_ratio is None: refineind=None else: raise ValueError("Adaptive tokenization is not set up/tested yet in TurbT") - + if field_labels_out is None: field_labels_out = state_labels + if self.diffusion: + emb = self.compute_diffusion_emb(sigma, imod) + + if diffusion_cond is not None and imod == self.nhlevels - 1: + diffusion_cond = rearrange(diffusion_cond, 'b t c d h w -> t b c d h w') + opts.diffusion_cond = diffusion_cond + if isgraph: """ For graph objects: support one level for now @@ -252,17 +266,26 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase): x = data if imodimod_bottom: opts.imod -= 1 x_pred = self.forward(x, state_labels, bcs, opts) #x_input = x.clone() #T,B,C,D,H,W T, B, _, D, H, W = x.shape - #self.debug_nan(x, message="input") - x, data_mean, data_std = normalize_spatiotemporal_persample(x) - #self.debug_nan(x, message="input after normalization") + if persample_normalize: + #self.debug_nan(x, message="input") + x, data_mean, data_std = normalize_spatiotemporal_persample(x) + #self.debug_nan(x, message="input after normalization") ################################################################################ if self.leadtime and leadtime is not None: leadtime = self.ltimeMLP[imod](leadtime) @@ -274,6 +297,21 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase): x, patch_ids, patch_ids_ref, mask_padding, _, _, tposarea_padding, _ = self.get_patchsequence(x, state_labels, tkhead_name, refineind=refineind, blockdict=blockdict, ilevel=imod, isgraph=isgraph) x = rearrange(x, 't b c ntoken_tot -> b c (t ntoken_tot)') ################################################################################ + if diffusion_cond is not None: + diffusion_cond_tokens, _, _, _, _, _, _, _ = \ + self.get_patchsequence(diffusion_cond, state_labels, tkhead_name, + refineind=None, blockdict=blockdict, + ilevel=imod, isgraph=isgraph) + diffusion_cond_tokens = rearrange( + diffusion_cond_tokens, 't b c ntoken_tot -> b c (t ntoken_tot)') + b_curr, _, L_curr = x.shape + cond_combined = diffusion_cond_tokens + emb.unsqueeze(-1) + cond_fused = self.diffusion_cond_proj[str(imod)]( + rearrange(cond_combined, 'b c L -> (b L) c')) + x = x + rearrange(cond_fused, '(b L) c -> b c L', b=b_curr) + elif self.diffusion: + x = x + emb.unsqueeze(-1) + ################################################################################ if self.posbias[imod] is not None and tposarea_padding is not None: use_zpos=True if D>1 else False posbias = self.posbias[imod](tposarea_padding, mask_padding=mask_padding, use_zpos=use_zpos) #b t L c->b t L c_emb @@ -345,7 +383,10 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase): #x_correct[:,var_index,...] = x_pred + x_correct[:,var_index,...] x_correct = x_correct + x_pred if imod==self.nhlevels-1: - x_correct=x_correct[:,state_labels[0],...] * data_std[-1] + data_mean[-1] + if persample_normalize: + x_correct=x_correct[:,state_labels[0],...] * data_std[-1] + data_mean[-1] + else: + x_correct=x_correct[:,state_labels[0],...] #since no T dim: b c d h w #x_correct=x_correct[:,var_index,...] * data_std[-1] + data_mean[-1] return x_correct #B,C_all,D,H,W for imod t b c d h w') + if isgraph: with torch.no_grad(): check_same_sample_across_halo(data, ghost_info, sequence_parallel_group) if sequence_parallel_group is not None else None @@ -158,9 +170,8 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase, train_opts: refineind=None else: refineind = get_top_variance_patchids(self.tokenizer_heads_params[tkhead_name], x, self.tokenizer_heads_gammaref[tkhead_name], refine_ratio) - #self.debug_nan(x, message="input") - x, data_mean, data_std = normalize_spatiotemporal_persample(x) - #self.debug_nan(x, message="input after normalization") + if persample_normalize: + x, data_mean, data_std = normalize_spatiotemporal_persample(x) ################################################################################ if self.leadtime and leadtime is not None: leadtime = self.ltimeMLP[imod](leadtime) @@ -179,6 +190,20 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase, train_opts: x_padding, patch_ids, patch_ids_ref, mask_padding, _, _, tposarea_padding, _ = self.get_patchsequence(x, state_labels, tkhead_name, refineind=refineind, blockdict=blockdict, ilevel=imod, isgraph=isgraph) x_padding = rearrange(x_padding, 't b c ntoken_tot -> b c (t ntoken_tot)') + if diffusion_cond is not None: + diffusion_cond_padding, _, _, _, _, _, _, _ = \ + self.get_patchsequence(diffusion_cond, state_labels, tkhead_name, + refineind=None, blockdict=blockdict, ilevel=imod, isgraph=isgraph) + diffusion_cond_padding = rearrange( + diffusion_cond_padding, 't b c ntoken_tot -> b c (t ntoken_tot)') + b_curr = x_padding.shape[0] + cond_combined = diffusion_cond_padding + emb.unsqueeze(-1) + cond_fused = self.diffusion_cond_proj[str(imod)]( + rearrange(cond_combined, 'b c L -> (b L) c')) + x_padding = x_padding + rearrange(cond_fused, '(b L) c -> b c L', b=b_curr) + elif self.diffusion: + x_padding = x_padding + emb.unsqueeze(-1) + # Repeat the steps for conditioning if present if conditioning: assert self.sts_model == False @@ -238,9 +263,11 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase, train_opts: ######### Denormalize ######## #t b c d h w x = x[:,:,field_labels_out[0],...] - x = x * data_std + data_mean + if persample_normalize: + x = x * data_std + data_mean ################################################################################ if train_opts is not None and train_opts.returnbase4train: - xbase = xbase * data_std + data_mean + if persample_normalize: + xbase = xbase * data_std + data_mean return x[-1], xbase[-1] return x[-1] diff --git a/matey/train.py b/matey/train.py index 6658b2e..b875895 100644 --- a/matey/train.py +++ b/matey/train.py @@ -21,6 +21,7 @@ from .models.svit import build_svit from .models.vit import build_vit from .models.turbt import build_turbt +from .models.diffusion_model import build_diffusion_model from .utils.logging_utils import Timer, record_function_opt from .utils.distributed_utils import get_sequence_parallel_group, add_weight_decay, CosineNoIncrease, determine_turt_levels from .utils.visualization_utils import checking_data_pred_tar @@ -33,13 +34,14 @@ import torch.distributed.checkpoint as dcp from torch.distributed.checkpoint.state_dict import get_state_dict, set_model_state_dict,set_optimizer_state_dict from torch_geometric.nn import global_mean_pool -from .utils.training_utils import autoregressive_rollout, compute_loss_and_logs, update_loss_logs_inplace_eval +from .utils.training_utils import EDMLoss, autoregressive_rollout, compute_loss_and_logs, update_loss_logs_inplace_eval import copy class Trainer: def __init__(self, params, global_rank, local_rank, device): self.device = device self.params = params + self.diffusion_config = getattr(params, 'diffusion_config', None) or {} self.global_rank = global_rank self.local_rank = local_rank self.world_size = int(os.environ.get("WORLD_SIZE", 1)) @@ -93,6 +95,9 @@ def __init__(self, params, global_rank, local_rank, device): print(f"Warning: reserved n_states {self.params.n_states} too small — {msg_suffix}.") self.params.n_states = new_n_states + if self.diffusion_config.get("diffusion", False): + self.diffusion_loss = EDMLoss() + self.initialize_model() self.initialize_optimizer() self.initialize_scheduler() @@ -179,7 +184,9 @@ def initialize_data(self): self.val_sampler.set_epoch(0) def initialize_model(self): - if self.params.model_type == 'avit': + if self.diffusion_config.get("diffusion", False): + self.model = build_diffusion_model(self.params).to(self.device) + elif self.params.model_type == 'avit': self.model = build_avit(self.params).to(self.device) elif self.params.model_type == "svit": self.model = build_svit(self.params).to(self.device) @@ -529,7 +536,10 @@ def set_field_dictionary(self): def model_forward(self, inp, field_labels, bcs, opts: ForwardOptionsBase, pushforward=True): # Handles a forward pass through the model, either normal or autoregressive rollout. autoregressive = getattr(self.params, "autoregressive", False) - if not autoregressive: + if self.diffusion_config.get("diffusion", False): + loss, output = self.diffusion_loss(self.model, inp, field_labels, bcs, opts) + return output, loss + elif not autoregressive: output = self.model(inp, field_labels, bcs, opts) return output, None else: @@ -546,8 +556,10 @@ def train_one_epoch(self): self.model.train() logs = {'train_rmse': torch.zeros(1).to(self.device), 'train_nrmse': torch.zeros(1).to(self.device), - 'train_l1': torch.zeros(1).to(self.device), - 'train_ssim': torch.zeros(1).to(self.device)} + 'train_l1': torch.zeros(1).to(self.device), + 'train_ssim': torch.zeros(1).to(self.device)} + if self.diffusion_config.get("diffusion", False): + logs['train_EDMloss'] = torch.zeros(1).to(self.device) steps = 0 grad_logs = defaultdict(lambda: torch.zeros(1, device=self.device)) grad_counts = defaultdict(lambda: torch.zeros(1, device=self.device)) @@ -631,11 +643,25 @@ def train_one_epoch(self): ghost_info = ghost_info if isgraph else None, ) with record_function_opt("model forward", enabled=self.profiling): - output, rollout_steps = self.model_forward(inp, field_labels, bcs, opts) - if not isgraph: - tar = tar[:, -1, :] #B,T(1 or leadtime),C,D,H,W -> B,C,D,H,W - #compute loss and update (in-place) logging dicts. - loss, log_nrmse = compute_loss_and_logs(output, tar, graphdata if isgraph else None, logs, loss_logs, dset_type, self.params) + if self.diffusion_config.get("diffusion", False): + if self.diffusion_config.get("cond_diffusion", False): + opts.diffusion_cond = rearrange(inp.to(self.device), 't b c d h w -> b t c d h w') + else: + assert torch.all(inp[0] == tar), "For unconditioned diffusion model, input should be the target" + inp_cond = rearrange(tar.to(self.device), 'b t c d h w -> t b c d h w') + output, loss_field = self.model_forward(inp_cond, field_labels, bcs, opts) + loss = loss_field.sum() / inp.shape[1] + print(f"Diffusion loss: {loss.item()}, EDM loss: {logs['train_EDMloss'].item()}") + else: + output, rollout_steps = self.model_forward(inp, field_labels, bcs, opts) + if not isgraph: + tar = tar[:, -1, :] #B,T(1 or leadtime),C,D,H,W -> B,C,D,H,W + if self.diffusion_config.get("diffusion", False): + logs['train_EDMloss'] += loss_field.mean().item() + _, log_nrmse = compute_loss_and_logs(output, tar, graphdata if isgraph else None, logs, loss_logs, dset_type, self.params) + else: + #compute loss and update (in-place) logging dicts. + loss, log_nrmse = compute_loss_and_logs(output, tar, graphdata if isgraph else None, logs, loss_logs, dset_type, self.params) bad = torch.isnan(loss).any() or torch.isinf(loss) torch.distributed.all_reduce(bad, op=torch.distributed.ReduceOp.SUM) if bad.item() > 0: @@ -719,6 +745,8 @@ def validate_one_epoch(self, full=False, cutoff_skip=False): 'valid_nrmse': torch.zeros(1).to(self.device), 'valid_l1': torch.zeros(1).to(self.device), 'valid_ssim': torch.zeros(1).to(self.device)} + if self.diffusion_config.get("diffusion", False): + logs['valid_EDMloss'] = torch.zeros(1).to(self.device) if cutoff_skip: return logs loss_dset_logs = {dataset.type: torch.zeros(1, device=self.device) for dataset in self.valid_dataset.sub_dsets} @@ -804,9 +832,18 @@ def validate_one_epoch(self, full=False, cutoff_skip=False): field_labels_out= field_labels_out, ghost_info = ghost_info if isgraph else None, ) - output, rollout_steps = self.model_forward(inp, field_labels, bcs, opts) - if not isgraph: - tar = tar[:, -1, :] #B,T(1 or leadtime),C,D,H,W -> B,C,D,H,W + if self.diffusion_config.get("diffusion", False): + if self.diffusion_config.get("cond_diffusion", False): + opts.diffusion_cond = rearrange(inp.to(self.device), 't b c d h w -> b t c d h w') + else: + assert torch.all(inp[0] == tar), "For unconditioned diffusion model, input should be the target" + inp_cond = rearrange(tar.to(self.device), 'b t c d h w -> t b c d h w') + output, loss_field = self.model_forward(inp_cond, field_labels, bcs, opts) + logs['valid_EDMloss'] += loss_field.mean().item() + else: + output, rollout_steps = self.model_forward(inp, field_labels, bcs, opts) + if not isgraph: + tar = tar[:, -1, :] #B,T(1 or leadtime),C,D,H,W -> B,C,D,H,W update_loss_logs_inplace_eval(output, tar, graphdata if isgraph else None, logs, loss_dset_logs, loss_l1_dset_logs, loss_rmse_dset_logs, dset_type) if not isgraph and getattr(self.params, "log_ssim", False): avg_ssim = get_ssim(output, tar, blockdict, self.global_rank, self.current_group, self.group_rank, self.group_size, self.device, self.valid_dataset, dset_index) diff --git a/matey/utils/forward_options.py b/matey/utils/forward_options.py index 4f292c9..9a181bb 100644 --- a/matey/utils/forward_options.py +++ b/matey/utils/forward_options.py @@ -26,4 +26,6 @@ class ForwardOptionsBase: #adaptive tokenization (1 of 2 settings) refine_ratio: Optional[float] = None imod_bottom: int = 0 #needed only by turbt - ghost_info: Optional = None \ No newline at end of file + ghost_info: Optional = None + sigma: Optional[Tensor] = None #needed only by diffusion model + diffusion_cond: Optional[Tensor] = None #needed only by conditional diffusion model \ No newline at end of file diff --git a/matey/utils/training_utils.py b/matey/utils/training_utils.py index d8f85cd..130f160 100644 --- a/matey/utils/training_utils.py +++ b/matey/utils/training_utils.py @@ -209,3 +209,26 @@ def update_loss_logs_inplace_eval(output, tar, graphdata, logs, loss_dset_logs, loss_l1_dset_logs[dset_type] += raw_l1_loss loss_rmse_dset_logs[dset_type] += raw_rmse_loss return + +class EDMLoss: + def __init__(self, P_mean=-1.2, P_std=1.2, sigma_data=0.5): + self.P_mean = P_mean + self.P_std = P_std + self.sigma_data = sigma_data + + def __call__(self, net, images, field_labels, bcs, opts, augment_pipe=None): + rnd_normal = torch.randn([1, images.shape[1], 1, 1, 1, 1], device=images.device) + sigma = (rnd_normal * self.P_std + self.P_mean).exp() + weight = (sigma ** 2 + self.sigma_data ** 2) / (sigma * self.sigma_data) ** 2 + # y, augment_labels = augment_pipe(images) if augment_pipe is not None else (images, None) + y = images + augment_labels = None + n = torch.randn_like(y) * sigma + # print(f"In EDMLoss: sigma shape: {sigma.shape}, y shape: {y.shape}, n shape: {n.shape}, labels shape: {opts.diffusion_cond.shape if getattr(opts, 'diffusion_cond', None) is not None else None}") + D_yn = net(y + n, sigma, field_labels, bcs, opts, augment_labels=augment_labels) + y = y.squeeze(0) if y.ndim == 6 else y + D_yn = D_yn.squeeze(0) if D_yn.ndim == 6 else D_yn + # print(f"D_yn shape: {D_yn.shape}, y shape: {y.shape}") + loss = weight * ((D_yn - y) ** 2) + return loss, D_yn +