From e349da55866235511bd4a4dc1c8dac8c4a12f3e6 Mon Sep 17 00:00:00 2001 From: Paul Laiu Date: Wed, 20 May 2026 10:44:17 -0400 Subject: [PATCH 01/12] adding diffusion model wrapper and four denoisers (turbt, turbt_modified, vit, unet) --- matey/models/__init__.py | 4 +- matey/models/diffusion_model.py | 82 ++++++ matey/models/time_modules.py | 51 ++++ matey/models/turbt.py | 119 ++++++++- matey/models/turbt_modified.py | 454 ++++++++++++++++++++++++++++++++ matey/models/unet.py | 354 +++++++++++++++++++++++++ matey/models/vit.py | 79 +++++- 7 files changed, 1134 insertions(+), 9 deletions(-) create mode 100644 matey/models/diffusion_model.py create mode 100644 matey/models/turbt_modified.py create mode 100644 matey/models/unet.py diff --git a/matey/models/__init__.py b/matey/models/__init__.py index fae8905..eae28a6 100644 --- a/matey/models/__init__.py +++ b/matey/models/__init__.py @@ -6,5 +6,5 @@ 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"] +from .unet import build_unet, UNet +__all__ = ["build_avit", "build_svit", "build_vit","build_turbt", "AViT","sViT_all2all","ViT_all2all","TurbT", "build_unet", "UNet"] diff --git a/matey/models/diffusion_model.py b/matey/models/diffusion_model.py new file mode 100644 index 0000000..cd812d4 --- /dev/null +++ b/matey/models/diffusion_model.py @@ -0,0 +1,82 @@ +from matey.models.unet import build_unet +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 +from .turbt_modified import build_turbt_modified +from .unet import build_unet + +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 + elif params.model_type == 'turbt_modified': + self.model = build_turbt_modified(params) + self.tokenizer_heads_params = self.model.tokenizer_heads_params + elif params.model_type == 'unet': + self.model = build_unet(params) + self.tokenizer_heads_params = {'CIFAR10': [[1, 1, 1]]} + 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/time_modules.py b/matey/models/time_modules.py index c903c81..52c74b4 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(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(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(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 69ab547..3b57edc 100644 --- a/matey/models/turbt.py +++ b/matey/models/turbt.py @@ -12,6 +12,8 @@ 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 +from .time_modules import FourierEmbedding, PositionalEmbedding, Linear +from torch.nn.functional import silu import sys, copy from operator import mul from functools import reduce @@ -41,7 +43,8 @@ def build_turbt(params): bias_type=params.bias_type, replace_patch=getattr(params, 'replace_patch', True), hierarchical=getattr(params, 'hierarchical', None), - notransposed=getattr(params, 'notransposed', False) + notransposed=getattr(params, 'notransposed', False), + diffusion=getattr(params, 'diffusion', False) ) return model @@ -58,7 +61,7 @@ class TurbT(BaseModel): sts_f """ def __init__(self, tokenizer_heads=None, embed_dim=768, num_heads=12, processor_blocks=8, n_states=6, - 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): + 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=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) self.drop_path = drop_path @@ -117,6 +120,33 @@ def __init__(self, tokenizer_heads=None, embed_dim=768, num_heads=12, processor 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)]) + self.diffusion=diffusion + if self.diffusion: + #from https://github.com/NVlabs/edm/blob/008a4e5316c8e3bfe61a62f874bddba254295afb/training/networks.py#L269 + #FIXME: @Paul, currently place holder pls check the proper setting of these variables + model_channels = 128 # Base multiplier for the number of channels. + channel_mult_emb = 4 # Multiplier for the dimensionality of the embedding vector. + embedding_type = 'positional' # Timestep embedding type: 'positional' for DDPM++, 'fourier' for NCSN++. + channel_mult_noise = 1 # Timestep embedding size: 1 for DDPM++, 2 for NCSN++. + emb_channels = model_channels * channel_mult_emb + noise_channels = model_channels * channel_mult_noise + init = dict(init_mode='xavier_uniform') + + self.map_noise = PositionalEmbedding(num_channels=noise_channels, endpoint=True) if embedding_type == 'positional' else FourierEmbedding(num_channels=noise_channels) + self.map_layer0 = Linear(in_features=noise_channels, out_features=emb_channels, **init) + # self.map_layer1 = Linear(in_features=emb_channels, out_features=embed_dim, **init) + self.map_layer1 = Linear(in_features=emb_channels, out_features=emb_channels, **init) + + self.affine = nn.ModuleDict({}) + for imod in range(self.nhlevels): + self.affine[str(imod)] = Linear(in_features=emb_channels, out_features=embed_dim, **init) + # self.affine = Linear(in_features=emb_channels, out_features=embed_dim, **init) + + self.skip_projection = nn.ModuleDict({}) + for imod in range(self.nhlevels): + self.skip_projection[str(imod)] = Linear(in_features=embed_dim*2, out_features=embed_dim, **init) + + 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}" @@ -222,6 +252,9 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase): cond_input = opts.cond_input isgraph=opts.isgraph field_labels_out=opts.field_labels_out + sigma = getattr(opts, 'sigma', None) + diffusion_cond = getattr(opts, 'diffusion_cond', None) + persample_normalize = False ################################################################## if refine_ratio is None: refineind=None @@ -231,6 +264,23 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase): if field_labels_out is None: field_labels_out = state_labels + if self.diffusion: + emb = self.map_noise(sigma) + # print(f'noise_labels shape: {sigma.shape}, emb shape: {emb.shape}') + emb = emb.reshape(emb.shape[0], 2, -1).flip(1).reshape(*emb.shape) # swap sin/cos + # print(f'after swap emb shape: {emb.shape}') + emb = silu(self.map_layer0(emb)) + # print(f'after first layer emb shape: {emb.shape}') + emb = silu(self.map_layer1(emb)) #in shape + # print(f'after second layer emb shape: {emb.shape}') + emb = self.affine[str(imod)](emb) + # print(f'after affine emb shape: {emb.shape}') + if diffusion_cond is not None: + if 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 + # print(f"Paul debugging, imod = {imod}, diffusion_cond shape after rearrange: {diffusion_cond.shape}", flush=True) + if isgraph: """ For graph objects: support one level for now @@ -247,8 +297,18 @@ 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) @@ -256,7 +316,8 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase): #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) + if persample_normalize: + 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: @@ -266,8 +327,20 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase): if self.cond_input and cond_input is not None: leadtime = self.inconMLP[imod](cond_input) if leadtime is None else leadtime+self.inconMLP[imod](cond_input) ########Encode and get patch sequences [B, C_emb, T*ntoken_len_tot]######## + # print(f"Paul debugging, diffusion_cond shape before get_patchsequence: {diffusion_cond.shape}, x shape: {x.shape}", flush=True) 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) + # print(f"Paul debugging, x shape after get_patchsequence: {x.shape}", flush=True) x = rearrange(x, 't b c ntoken_tot -> b c (t ntoken_tot)') + # print(f"Paul debugging, x shape after rearrange: {x.shape}", flush=True) + if diffusion_cond is not None: + diffusion_cond, _, _, _, _, _, _, _ = self.get_patchsequence(diffusion_cond, state_labels, tkhead_name, refineind=refineind, blockdict=blockdict, ilevel=imod, isgraph=isgraph) + # print(f"Paul debugging, diffusion_cond shape after get_patchsequence: {diffusion_cond.shape}", flush=True) + diffusion_cond = rearrange(diffusion_cond, 't b c ntoken_tot -> b c (t ntoken_tot)') + # print(f"Paul debugging, diffusion_cond shape after get_patchsequence reshaping: {diffusion_cond.shape}", flush=True) + + # if self.diffusion: + # print("Paul debugging, sigma, emb, x", sigma.shape, emb.shape, x.shape, flush=True) + # 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 @@ -287,16 +360,47 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase): ps = self.tokenizer_ensemble_heads[imod][tkhead_name]["embed"][-1].patch_size nfact=4//ps[-1] """ + if diffusion_cond is not None: + # print(f"Paul debugging, before sequence_factor_short, diffusion_cond shape: {diffusion_cond.shape}, x shape: {x.shape}", flush=True) + diffusion_cond, _=self.sequence_factor_short(diffusion_cond, imod, tkhead_name, [T, D, H, W], nfact=nfact) + # print(f"Paul debugging, after sequence_factor_short, diffusion_cond shape: {diffusion_cond.shape}, x shape: {x.shape}", flush=True) + x, nfact=self.sequence_factor_short(x, imod, tkhead_name, [T, D, H, W], nfact=nfact) + for iblk, blk in enumerate(self.module_blocks[str(imod)]): if iblk==0: b_mod=x.shape[0] if not isgraph and leadtime is not None: leadtime = leadtime.repeat(b_mod // B, 1) + if self.diffusion: + if diffusion_cond is not None: + # add diffusion_cond as additional tokens for the diffusion model to attend to + x = torch.cat([x, diffusion_cond], dim=2) + + # add noise embedding to input tokens as conditional information for diffusion model + emb_mod = emb.repeat(b_mod // B, 1) + x = torch.cat([x, emb_mod.unsqueeze(-1)], dim=2) + + x_input = x.clone() x = blk(x, sequence_parallel_group=sequence_parallel_group, bcs=bcs, leadtime=leadtime, mask_padding=mask4attblk, local_att=local_att) else: + if iblk==len(self.module_blocks[str(imod)])-1 and self.diffusion: + # for the last block, add skip connection from input tokens to facilitate training of diffusion model + x = torch.cat([x, x_input], dim=1) + local_batch = x.shape[0] + x = rearrange(x, 'b c L -> (b L) c') + x = self.skip_projection[str(imod)](x) # Residual connection from input to the last block in the level + x = rearrange(x, '(b L) c -> b c L', b = local_batch) x = blk(x, sequence_parallel_group=sequence_parallel_group, bcs=bcs, leadtime=None, mask_padding=mask4attblk, local_att=local_att) + #self.debug_nan(x_padding, message="attention block") + + if self.diffusion: + x = x[:, :, :-1] # remove noise embedding from tokens after processing + if diffusion_cond is not None: + doubled_L = x.shape[2] + x = x[:, :, :doubled_L//2] # remove diffusion_cond part from tokens after processing + if local_att: #nfact=4//ps[-1] x=self.sequence_factor_long(x, imod, tkhead_name, [T, D, H, W], nfact=nfact) @@ -337,7 +441,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= 0, + cond_input=getattr(params, 'supportdata', False), + n_steps=params.n_steps, + bias_type=params.bias_type, + replace_patch=getattr(params, 'replace_patch', True), + hierarchical=getattr(params, 'hierarchical', None), + notransposed=getattr(params, 'notransposed', False), + diffusion=getattr(params, 'diffusion', False), + ) + return model + + +class TurbTModified(BaseModel): + """ + TurbT with DiffiT-inspired diffusion conditioning to eliminate patch artifacts. + + Identical to TurbT for all non-diffusion code paths. + """ + + def __init__(self, tokenizer_heads=None, embed_dim=768, num_heads=12, + processor_blocks=8, n_states=6, 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=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, + ) + self.drop_path = drop_path + self.dp = np.linspace(0, drop_path, processor_blocks) + self.module_blocks = nn.ModuleDict({}) + self.sts_model = sts_model + self.sts_train = sts_train + + self.num_heads = num_heads + self.n_steps = n_steps + self.processor_blocks = processor_blocks + self.replace_patch = replace_patch + assert not (self.replace_patch and self.sts_model) + + self.upscale_factors = [1] + self.module_upscale = nn.ModuleDict({}) + self.module_upscale_space = nn.ModuleDict({}) + self.module_upscale_space2D = nn.ModuleDict({}) + + self.hierarchical = False + self.datafilter_kernel = None + if hierarchical is not None: + self.hierarchical = True + filtersize = hierarchical["filtersize"] + self.datafilter_kernel = construct_filterkernel(filtersize) + self.datafilter_kernel2D = construct_filterkernel2D(filtersize) + self.filtersize = filtersize + self.nhlevels = hierarchical["nlevels"] + self.upscale_factors = [1] + [self.filtersize for _ in range(self.nhlevels - 1)] + + for imod, upscalefactor in enumerate(self.upscale_factors): + if hierarchical["fixedupsample"]: + self.module_upscale_space[str(imod)] = nn.Upsample( + scale_factor=(upscalefactor, upscalefactor, upscalefactor), + mode='trilinear', align_corners=True) + self.module_upscale_space2D[str(imod)] = nn.Upsample( + scale_factor=(1, upscalefactor, upscalefactor), + mode='trilinear', align_corners=True) + elif hierarchical["linearupsample"]: + self.module_upscale_space[str(imod)] = torch.nn.Sequential( + nn.Upsample(scale_factor=(upscalefactor, upscalefactor, upscalefactor), + mode='trilinear', align_corners=True), + nn.Conv3d(n_states, n_states, + kernel_size=(upscalefactor, upscalefactor, upscalefactor), + stride=1, padding="same", bias=True, padding_mode="reflect"), + nn.InstanceNorm3d(n_states, affine=True)) + self.module_upscale_space2D[str(imod)] = torch.nn.Sequential( + nn.Upsample(scale_factor=(1, upscalefactor, upscalefactor), + mode='trilinear', align_corners=True), + nn.Conv3d(n_states, n_states, + kernel_size=(1, upscalefactor, upscalefactor), + stride=1, padding="same", bias=True, padding_mode="reflect"), + nn.InstanceNorm3d(n_states, affine=True)) + else: + self.module_upscale_space[str(imod)] = UpsampleinSpace( + patch_size=[upscalefactor, upscalefactor, upscalefactor], channels=n_states) + self.module_upscale_space2D[str(imod)] = UpsampleinSpace( + patch_size=[1, upscalefactor, upscalefactor], channels=n_states) + + if imod == 0: + 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)]) + else: + 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)]) + + self.diffusion = diffusion + if self.diffusion: + model_channels = 128 + channel_mult_emb = 4 + embedding_type = 'positional' + 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(num_channels=noise_channels, endpoint=True) + if embedding_type == 'positional' + else FourierEmbedding(num_channels=noise_channels)) + self.map_layer0 = Linear(in_features=noise_channels, out_features=emb_channels, **init) + self.map_layer1 = Linear(in_features=emb_channels, out_features=emb_channels, **init) + + self.affine = nn.ModuleDict({}) + for imod in range(self.nhlevels): + self.affine[str(imod)] = Linear(in_features=emb_channels, out_features=embed_dim, **init) + + # NEW: per-level learned projection for fusing diffusion_cond tokens into + # the main token stream via element-wise addition (replaces token appending). + self.diffusion_cond_proj = nn.ModuleDict({}) + for imod in range(self.nhlevels): + self.diffusion_cond_proj[str(imod)] = Linear(in_features=embed_dim, out_features=embed_dim, **init) + + # ------------------------------------------------------------------ + # Helpers copied verbatim from TurbT + # ------------------------------------------------------------------ + + def filterdata(self, data, blockdict=None): + assert data.ndim == 6, f"unkown tensor shape in filter_data, {data.shape}" + with torch.no_grad(): + kernel_size = self.filtersize + T, B, C, D, H, W = data.shape + data = rearrange(data, 't b c d h w -> (t b c) d h w') + if D == 1: + kernel = self.datafilter_kernel2D + filtered = F.conv3d(data[:, None, :, :, :], kernel.to(data.device), + stride=(1, kernel_size, kernel_size)) + else: + kernel = self.datafilter_kernel + filtered = F.conv3d(data[:, None, :, :, :], kernel.to(data.device), + stride=kernel_size) + filtered = rearrange(filtered, '(t b c) c1 d h w -> t b (c c1) d h w', t=T, b=B, c=C) + if blockdict is not None: + assert [D, H, W] == blockdict["Ind_dim"], f"(D,H,W),{(D,H,W)}, {blockdict['Ind_dim']}" + if D == 1: + blockdict["Ind_dim"] = [D, H // kernel_size, W // kernel_size] + else: + blockdict["Ind_dim"] = [D // kernel_size, H // kernel_size, W // kernel_size] + return filtered, blockdict + + def upsampeldata(self, data, imod): + B, C, D, H, W = data.shape + if D == 1: + data_upsample = self.module_upscale_space2D[str(imod)](data) + else: + data_upsample = self.module_upscale_space[str(imod)](data) + return data_upsample + + def sequence_factor_short(self, x, ilevel, tkhead_name, tspace_dims, nfact=2): + B, C, TL = x.shape + embed_ensemble = self.tokenizer_ensemble_heads[ilevel][tkhead_name]["embed"] + ntokendim = [] + ps_c = embed_ensemble[-1].patch_size + for idim, dim in enumerate(tspace_dims[1:]): + ntokendim.append(dim // ps_c[idim]) + assert TL == tspace_dims[0] * reduce(mul, ntokendim), \ + f"{TL}, {tspace_dims}, {ntokendim}" + d, h, w = ntokendim + if h // nfact < 4: + nfact = max(1, h // 4) + if nfact < 2: + return x, nfact + nfactd = 1 if d == 1 else nfact + x = rearrange(x, 'b c (t d h w) -> b c t d h w', + t=tspace_dims[0], d=d, h=h, w=w) + x = x.unfold(3, d // nfactd, d // nfactd).unfold(4, h // nfact, h // nfact).unfold(5, w // nfact, w // nfact) + x = rearrange(x, 'b c t nd nh nw d h w -> (b nd nh nw) c (t d h w)') + return x, nfact + + def sequence_factor_long(self, x, ilevel, tkhead_name, tspace_dims, nfact=2): + if nfact < 2: + return x + B, C, TL = x.shape + embed_ensemble = self.tokenizer_ensemble_heads[ilevel][tkhead_name]["embed"] + ntokendim = [] + ps_c = embed_ensemble[-1].patch_size + for idim, dim in enumerate(tspace_dims[1:]): + ntokendim.append(dim // ps_c[idim]) + d, h, w = ntokendim + nfactd = 1 if d == 1 else nfact + assert TL * (nfactd * nfact * nfact) == tspace_dims[0] * reduce(mul, ntokendim), \ + f"{TL}, {tspace_dims}, {ntokendim}, {nfact, nfactd}" + x = rearrange(x, '(b nd nh nw) c (t d h w) -> b c t nd nh nw d h w', + b=B // (nfactd * nfact * nfact), nd=nfactd, nh=nfact, nw=nfact, + d=d // nfactd, h=h // nfact, w=w // nfact) + x = rearrange(x, 'b c t nd nh nw d h w -> b c t (nd d) (nh h) (nw w)') + x = rearrange(x, 'b c t d h w -> b c (t d h w)') + return x + + # ------------------------------------------------------------------ + # Forward + # ------------------------------------------------------------------ + + def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase): + # ---- unpack arguments ------------------------------------------------ + imod = opts.imod + imod_bottom = opts.imod_bottom + tkhead_name = opts.tkhead_name + sequence_parallel_group = opts.sequence_parallel_group + leadtime = opts.leadtime + blockdict = opts.blockdict + refine_ratio = opts.refine_ratio + cond_input = opts.cond_input + isgraph = opts.isgraph + field_labels_out = opts.field_labels_out + sigma = getattr(opts, 'sigma', None) + diffusion_cond = getattr(opts, 'diffusion_cond', None) + persample_normalize = False + # ---- argument checks ------------------------------------------------- + if refine_ratio is not None: + raise ValueError("Adaptive tokenization is not set up/tested yet in TurbTModified") + + if field_labels_out is None: + field_labels_out = state_labels + + # ---- noise embedding (diffusion only) -------------------------------- + if self.diffusion: + emb = self.map_noise(sigma) + emb = emb.reshape(emb.shape[0], 2, -1).flip(1).reshape(*emb.shape) # swap sin/cos + emb = silu(self.map_layer0(emb)) + emb = silu(self.map_layer1(emb)) + emb = self.affine[str(imod)](emb) # (B, embed_dim) + + # Rearrange diffusion_cond to (T, B, C, D, H, W) at finest level. + # At coarser levels it arrives pre-filtered in the right format. + 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 + + # ---- graph path (unchanged from TurbT) -------------------------------- + if isgraph: + x = data.x + edge_index = data.edge_index + batch = data.batch + T = x.shape[1] + x, data_mean, data_std = normalize_spatiotemporal_persample_graph(x, batch) + refineind = None + x = (x, batch, edge_index) + else: + x = data + + if imod < self.nhlevels - 1: + # Filter x and diffusion_cond together so they share the same + # coarse grid (unchanged from TurbT). + if diffusion_cond is not None: + concat_data = torch.cat([x, diffusion_cond], dim=2) + concat_data, blockdict = self.filterdata(concat_data, blockdict=blockdict) + x = concat_data[:, :, :x.shape[2], :, :, :] + diffusion_cond = concat_data[:, :, x.shape[2]:, :, :, :] + opts.diffusion_cond = diffusion_cond + opts.blockdict = blockdict + else: + x, blockdict = self.filterdata(x, blockdict=blockdict) + opts.blockdict = blockdict + + if imod > imod_bottom: + opts.imod -= 1 + x_pred = self.forward(x, state_labels, bcs, opts) + + T, B, _, D, H, W = x.shape + if persample_normalize: + x, data_mean, data_std = normalize_spatiotemporal_persample(x) + + # ---- leadtime / conditional input (unchanged) ------------------------- + if self.leadtime and leadtime is not None: + leadtime = self.ltimeMLP[imod](leadtime) + else: + leadtime = None + if self.cond_input and cond_input is not None: + leadtime = (self.inconMLP[imod](cond_input) + if leadtime is None + else leadtime + self.inconMLP[imod](cond_input)) + + # ---- tokenize x ------------------------------------------------------- + x, patch_ids, patch_ids_ref, mask_padding, _, _, tposarea_padding, _ = \ + self.get_patchsequence(x, state_labels, tkhead_name, refineind=None, + blockdict=blockdict, ilevel=imod, isgraph=isgraph) + x = rearrange(x, 't b c ntoken_tot -> b c (t ntoken_tot)') + + # ---- fuse noise embedding and diffusion_cond into x (CHANGED) -------- + # Both the noise embedding (emb) and diffusion_cond are combined and + # added into x before the block loop. When diffusion_cond is present, + # emb is added to the cond tokens before projection so both signals are + # fused in one learned step. When diffusion_cond is absent, emb is + # broadcast directly into x. + 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) + + # ---- positional bias (unchanged) -------------------------------------- + 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) + posbias = rearrange(posbias, 'b t L c -> b c (t L)') + x = x + posbias + del posbias + + # ---- attention mask --------------------------------------------------- + mask4attblk = None if (mask_padding is not None and mask_padding.all()) else mask_padding + + # ---- local attention windowing (CHANGED: no separate cond handling) --- + # diffusion_cond has already been fused into x, so only x needs windowing. + local_att = not isgraph and imod > imod_bottom + if local_att: + nfact = (max(2 ** (2 * (imod - imod_bottom)) // blockdict["nproc_blocks"][-1], 1) + if blockdict is not None + else max(2 ** (2 * (imod - imod_bottom)), 1)) + x, nfact = self.sequence_factor_short(x, imod, tkhead_name, [T, D, H, W], nfact=nfact) + + # ---- transformer block loop (CHANGED) --------------------------------- + # Key changes: + # (a) No time embedding token appended to the sequence. + # (b) Noise embedding is fused into x before the loop (not via leadtime). + # (c) leadtime passed only at iblk==0, None thereafter (original behaviour). + for iblk, blk in enumerate(self.module_blocks[str(imod)]): + if iblk == 0: + b_mod = x.shape[0] + if not isgraph and leadtime is not None: + leadtime = leadtime.repeat(b_mod // B, 1) + + block_cond = leadtime if iblk == 0 else None + + x = blk(x, sequence_parallel_group=sequence_parallel_group, + bcs=bcs, leadtime=block_cond, mask_padding=mask4attblk, + local_att=local_att) + + # No time/cond tokens to strip — they were never appended. + + # ---- un-window (unchanged) -------------------------------------------- + if local_att: + x = self.sequence_factor_long(x, imod, tkhead_name, [T, D, H, W], nfact=nfact) + + # ---- decode ----------------------------------------------------------- + x = rearrange(x, 'b c (t ntoken_tot) -> t b c ntoken_tot', t=T) + + if isgraph: + x = rearrange(x, 't b c ntoken_tot -> b ntoken_tot t c') + x = densenodes_to_graphnodes(x, mask_padding) + x = (x, batch, edge_index) + D, H, W = -1, -1, -1 + + x = self.get_spatiotemporalfromsequence( + x, patch_ids, patch_ids_ref, [D, H, W], tkhead_name, + ilevel=imod, isgraph=isgraph) + + if isgraph: + node_ft, batch, edge_index = x + x = node_ft[:, :, field_labels_out[0]] + N = x.shape[0] + mask = torch.isin(state_labels[0], field_labels_out[0]) + mean_node = data_mean[batch].view(N, 1, -1)[:, :, mask] + std_node = data_std[batch].view(N, 1, -1)[:, :, mask] + x = x * std_node + mean_node + return x[:, -1, :] + + # ---- hierarchical upsampling (unchanged) ------------------------------ + x_correct = x[-1] + del x + if imod > imod_bottom: + x_filter = self.filterdata(x_correct[None, ...])[0][-1] + filtered_eps = self.upsampeldata(x_filter, imod) + x_correct = x_correct - filtered_eps + x_pred = self.upsampeldata(x_pred, imod) + x_correct = x_correct + x_pred + + if imod == self.nhlevels - 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], ...] + + return x_correct diff --git a/matey/models/unet.py b/matey/models/unet.py new file mode 100644 index 0000000..c2659f5 --- /dev/null +++ b/matey/models/unet.py @@ -0,0 +1,354 @@ +# SPDX-License-Identifier: MIT +# SPDX-FileCopyrightText: 2026 UT-Battelle, LLC +# This file is part of the MATEY Project. + +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +from einops import rearrange +from .spacetime_modules import SpaceTimeBlock_all2all +from .basemodel import BaseModel +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 +from .time_modules import FourierEmbedding, PositionalEmbedding, Linear +from torch.nn.functional import silu +import sys, copy +from operator import mul +from functools import reduce +from ..utils.forward_options import ForwardOptionsBase +import torch.distributed as dist +from ..utils import densenodes_to_graphnodes + + +def build_unet(params): + """ Builds model from parameter file. + """ + model = UNet() + return model + +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(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 + + +class Conv2d(torch.nn.Module): + def __init__(self, + in_channels, out_channels, kernel, bias=True, up=False, down=False, + resample_filter=[1,1], fused_resample=False, init_mode='kaiming_normal', init_weight=1, init_bias=0, + ): + assert not (up and down) + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.up = up + self.down = down + self.fused_resample = fused_resample + init_kwargs = dict(mode=init_mode, fan_in=in_channels*kernel*kernel, fan_out=out_channels*kernel*kernel) + self.weight = torch.nn.Parameter(weight_init([out_channels, in_channels, kernel, kernel], **init_kwargs) * init_weight) if kernel else None + self.bias = torch.nn.Parameter(weight_init([out_channels], **init_kwargs) * init_bias) if kernel and bias else None + f = torch.as_tensor(resample_filter, dtype=torch.float32) + f = f.ger(f).unsqueeze(0).unsqueeze(1) / f.sum().square() + self.register_buffer('resample_filter', f if up or down else None) + + def forward(self, x): + w = self.weight.to(x.dtype) if self.weight is not None else None + b = self.bias.to(x.dtype) if self.bias is not None else None + f = self.resample_filter.to(x.dtype) if self.resample_filter is not None else None + w_pad = w.shape[-1] // 2 if w is not None else 0 + f_pad = (f.shape[-1] - 1) // 2 if f is not None else 0 + + if self.fused_resample and self.up and w is not None: + x = torch.nn.functional.conv_transpose2d(x, f.mul(4).tile([self.in_channels, 1, 1, 1]), groups=self.in_channels, stride=2, padding=max(f_pad - w_pad, 0)) + x = torch.nn.functional.conv2d(x, w, padding=max(w_pad - f_pad, 0)) + elif self.fused_resample and self.down and w is not None: + x = torch.nn.functional.conv2d(x, w, padding=w_pad+f_pad) + x = torch.nn.functional.conv2d(x, f.tile([self.out_channels, 1, 1, 1]), groups=self.out_channels, stride=2) + else: + if self.up: + x = torch.nn.functional.conv_transpose2d(x, f.mul(4).tile([self.in_channels, 1, 1, 1]), groups=self.in_channels, stride=2, padding=f_pad) + if self.down: + x = torch.nn.functional.conv2d(x, f.tile([self.in_channels, 1, 1, 1]), groups=self.in_channels, stride=2, padding=f_pad) + if w is not None: + x = torch.nn.functional.conv2d(x, w, padding=w_pad) + if b is not None: + x = x.add_(b.reshape(1, -1, 1, 1)) + return x + + +class GroupNorm(torch.nn.Module): + def __init__(self, num_channels, num_groups=32, min_channels_per_group=4, eps=1e-5): + super().__init__() + self.num_groups = min(num_groups, num_channels // min_channels_per_group) + self.eps = eps + self.weight = torch.nn.Parameter(torch.ones(num_channels)) + self.bias = torch.nn.Parameter(torch.zeros(num_channels)) + + def forward(self, x): + x = torch.nn.functional.group_norm(x, num_groups=self.num_groups, weight=self.weight.to(x.dtype), bias=self.bias.to(x.dtype), eps=self.eps) + return x + +#---------------------------------------------------------------------------- +# Attention weight computation, i.e., softmax(Q^T * K). +# Performs all computation using FP32, but uses the original datatype for +# inputs/outputs/gradients to conserve memory. + +class AttentionOp(torch.autograd.Function): + @staticmethod + def forward(ctx, q, k): + w = torch.einsum('ncq,nck->nqk', q.to(torch.float32), (k / np.sqrt(k.shape[1])).to(torch.float32)).softmax(dim=2).to(q.dtype) + ctx.save_for_backward(q, k, w) + return w + + @staticmethod + def backward(ctx, dw): + q, k, w = ctx.saved_tensors + db = torch._softmax_backward_data(grad_output=dw.to(torch.float32), output=w.to(torch.float32), dim=2, input_dtype=torch.float32) + dq = torch.einsum('nck,nqk->ncq', k.to(torch.float32), db).to(q.dtype) / np.sqrt(k.shape[1]) + dk = torch.einsum('ncq,nqk->nck', q.to(torch.float32), db).to(k.dtype) / np.sqrt(k.shape[1]) + return dq, dk + + +class UNetBlock(torch.nn.Module): + def __init__(self, + in_channels, out_channels, emb_channels, up=False, down=False, attention=False, + num_heads=None, channels_per_head=64, dropout=0, skip_scale=1, eps=1e-5, + resample_filter=[1,1], resample_proj=False, adaptive_scale=True, + init=dict(), init_zero=dict(init_weight=0), init_attn=None, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.emb_channels = emb_channels + self.num_heads = 0 if not attention else num_heads if num_heads is not None else out_channels // channels_per_head + self.dropout = dropout + self.skip_scale = skip_scale + self.adaptive_scale = adaptive_scale + + self.norm0 = GroupNorm(num_channels=in_channels, eps=eps) + self.conv0 = Conv2d(in_channels=in_channels, out_channels=out_channels, kernel=3, up=up, down=down, resample_filter=resample_filter, **init) + self.affine = Linear(in_features=emb_channels, out_features=out_channels*(2 if adaptive_scale else 1), **init) + self.norm1 = GroupNorm(num_channels=out_channels, eps=eps) + self.conv1 = Conv2d(in_channels=out_channels, out_channels=out_channels, kernel=3, **init_zero) + + self.skip = None + if out_channels != in_channels or up or down: + kernel = 1 if resample_proj or out_channels!= in_channels else 0 + self.skip = Conv2d(in_channels=in_channels, out_channels=out_channels, kernel=kernel, up=up, down=down, resample_filter=resample_filter, **init) + + if self.num_heads: + self.norm2 = GroupNorm(num_channels=out_channels, eps=eps) + self.qkv = Conv2d(in_channels=out_channels, out_channels=out_channels*3, kernel=1, **(init_attn if init_attn is not None else init)) + self.proj = Conv2d(in_channels=out_channels, out_channels=out_channels, kernel=1, **init_zero) + + def forward(self, x, emb): + orig = x + x = self.conv0(silu(self.norm0(x))) + + params = self.affine(emb).unsqueeze(2).unsqueeze(3).to(x.dtype) + + # print(f'emb shape: {emb.shape}, affine out shape: {self.affine(emb).shape}, x shape: {x.shape}, params shape: {params.shape}') + + if self.adaptive_scale: + scale, shift = params.chunk(chunks=2, dim=1) + x = silu(torch.addcmul(shift, self.norm1(x), scale + 1)) + else: + x = silu(self.norm1(x.add_(params))) + + x = self.conv1(torch.nn.functional.dropout(x, p=self.dropout, training=self.training)) + x = x.add_(self.skip(orig) if self.skip is not None else orig) + x = x * self.skip_scale + + if self.num_heads: + q, k, v = self.qkv(self.norm2(x)).reshape(x.shape[0] * self.num_heads, x.shape[1] // self.num_heads, 3, -1).unbind(2) + w = AttentionOp.apply(q, k) + a = torch.einsum('nqk,nck->ncq', w, v) + x = self.proj(a.reshape(*x.shape)).add_(x) + x = x * self.skip_scale + return x + + +class UNet(torch.nn.Module): + def __init__(self, + img_resolution = 256, # Image resolution at input/output. + in_channels = 8, # Number of color channels at input. + out_channels = 8, # Number of color channels at output. + label_dim = 0, # Number of class labels, 0 = unconditional. + augment_dim = 0, # Augmentation label dimensionality, 0 = no augmentation. + + model_channels = 128, # Base multiplier for the number of channels. + channel_mult = [1,2,2,2], # Per-resolution multipliers for the number of channels. + channel_mult_emb = 4, # Multiplier for the dimensionality of the embedding vector. + num_blocks = 4, # Number of residual blocks per resolution. + attn_resolutions = [16], # List of resolutions with self-attention. + dropout = 0.10, # Dropout probability of intermediate activations. + label_dropout = 0, # Dropout probability of class labels for classifier-free guidance. + + embedding_type = 'positional', # Timestep embedding type: 'positional' for DDPM++, 'fourier' for NCSN++. + channel_mult_noise = 1, # Timestep embedding size: 1 for DDPM++, 2 for NCSN++. + encoder_type = 'standard', # Encoder architecture: 'standard' for DDPM++, 'residual' for NCSN++. + decoder_type = 'standard', # Decoder architecture: 'standard' for both DDPM++ and NCSN++. + resample_filter = [1,1], # Resampling filter: [1,1] for DDPM++, [1,3,3,1] for NCSN++. + ): + assert embedding_type in ['fourier', 'positional'] + assert encoder_type in ['standard', 'skip', 'residual'] + assert decoder_type in ['standard', 'skip'] + + super().__init__() + self.label_dropout = label_dropout + emb_channels = model_channels * channel_mult_emb + noise_channels = model_channels * channel_mult_noise + init = dict(init_mode='xavier_uniform') + init_zero = dict(init_mode='xavier_uniform', init_weight=1e-5) + init_attn = dict(init_mode='xavier_uniform', init_weight=np.sqrt(0.2)) + block_kwargs = dict( + emb_channels=emb_channels, num_heads=1, dropout=dropout, skip_scale=np.sqrt(0.5), eps=1e-6, + resample_filter=resample_filter, resample_proj=True, adaptive_scale=False, + init=init, init_zero=init_zero, init_attn=init_attn, + ) + + # Mapping. + self.map_noise = PositionalEmbedding(num_channels=noise_channels, endpoint=True) if embedding_type == 'positional' else FourierEmbedding(num_channels=noise_channels) + self.map_label = Linear(in_features=label_dim, out_features=noise_channels, **init) if label_dim else None + self.map_augment = Linear(in_features=augment_dim, out_features=noise_channels, bias=False, **init) if augment_dim else None + self.map_layer0 = Linear(in_features=noise_channels, out_features=emb_channels, **init) + self.map_layer1 = Linear(in_features=emb_channels, out_features=emb_channels, **init) + + # print(f"model_channels: {model_channels}, noise_channels: {noise_channels}, emb_channels: {emb_channels}") + # Encoder. + self.enc = torch.nn.ModuleDict() + cout = in_channels + caux = in_channels + for level, mult in enumerate(channel_mult): + res = img_resolution >> level + if level == 0: + cin = cout + cout = model_channels + self.enc[f'{res}x{res}_conv'] = Conv2d(in_channels=cin, out_channels=cout, kernel=3, **init) + else: + self.enc[f'{res}x{res}_down'] = UNetBlock(in_channels=cout, out_channels=cout, down=True, **block_kwargs) + if encoder_type == 'skip': + self.enc[f'{res}x{res}_aux_down'] = Conv2d(in_channels=caux, out_channels=caux, kernel=0, down=True, resample_filter=resample_filter) + self.enc[f'{res}x{res}_aux_skip'] = Conv2d(in_channels=caux, out_channels=cout, kernel=1, **init) + if encoder_type == 'residual': + self.enc[f'{res}x{res}_aux_residual'] = Conv2d(in_channels=caux, out_channels=cout, kernel=3, down=True, resample_filter=resample_filter, fused_resample=True, **init) + caux = cout + for idx in range(num_blocks): + cin = cout + cout = model_channels * mult + attn = (res in attn_resolutions) + self.enc[f'{res}x{res}_block{idx}'] = UNetBlock(in_channels=cin, out_channels=cout, attention=attn, **block_kwargs) + skips = [block.out_channels for name, block in self.enc.items() if 'aux' not in name] + + # Decoder. + self.dec = torch.nn.ModuleDict() + for level, mult in reversed(list(enumerate(channel_mult))): + res = img_resolution >> level + if level == len(channel_mult) - 1: + self.dec[f'{res}x{res}_in0'] = UNetBlock(in_channels=cout, out_channels=cout, attention=True, **block_kwargs) + self.dec[f'{res}x{res}_in1'] = UNetBlock(in_channels=cout, out_channels=cout, **block_kwargs) + else: + self.dec[f'{res}x{res}_up'] = UNetBlock(in_channels=cout, out_channels=cout, up=True, **block_kwargs) + for idx in range(num_blocks + 1): + cin = cout + skips.pop() + cout = model_channels * mult + attn = (idx == num_blocks and res in attn_resolutions) + self.dec[f'{res}x{res}_block{idx}'] = UNetBlock(in_channels=cin, out_channels=cout, attention=attn, **block_kwargs) + if decoder_type == 'skip' or level == 0: + if decoder_type == 'skip' and level < len(channel_mult) - 1: + self.dec[f'{res}x{res}_aux_up'] = Conv2d(in_channels=out_channels, out_channels=out_channels, kernel=0, up=True, resample_filter=resample_filter) + self.dec[f'{res}x{res}_aux_norm'] = GroupNorm(num_channels=cout, eps=1e-6) + self.dec[f'{res}x{res}_aux_conv'] = Conv2d(in_channels=cout, out_channels=out_channels, kernel=3, **init_zero) + + # def forward(self, x, noise_labels, class_labels, augment_labels=None): + def forward(self, x, field_labels, bcs, opts, augment_labels=None): + # Mapping. + # print(f'Input x shape: {x.shape}') + x = x.squeeze(0) + x = x.squeeze(2) # reshape from [1, B, C, 1, H, W] to [B, C, H, W] + # print(f'Reshaped x shape: {x.shape}') + noise_labels = getattr(opts, 'sigma', None) + diffusion_cond = getattr(opts, 'diffusion_cond', None) + if diffusion_cond is not None: + diffusion_cond = diffusion_cond.squeeze(1) + diffusion_cond = diffusion_cond.squeeze(2) # reshape from [B, 1, C, 1, H, W] to [B, C, H, W] + # print(f'Reshaped diffusion_cond shape: {diffusion_cond.shape}') + x = torch.cat([x, diffusion_cond], dim=1) # concatenate along channel dimension to get shape [B, 2*C, H, W] + # print(f'concatenated x shape: {x.shape}') + + if noise_labels is not None: + emb = self.map_noise(noise_labels) + emb = emb.reshape(emb.shape[0], 2, -1).flip(1).reshape(*emb.shape) # swap sin/cos + # if self.map_label is not None: + # class_labels = opts.diffusion_cond.reshape(opts.diffusion_cond.shape[0], -1) if hasattr(opts, 'diffusion_cond') else None + # tmp = class_labels + # if self.training and self.label_dropout: + # tmp = tmp * (torch.rand([x.shape[0], 1], device=x.device) >= self.label_dropout).to(tmp.dtype) + # emb = emb + self.map_label(tmp * np.sqrt(self.map_label.in_features)) + # if self.map_augment is not None and augment_labels is not None: + # emb = emb + self.map_augment(augment_labels) + + emb = silu(self.map_layer0(emb)) + emb = silu(self.map_layer1(emb)) + + # Encoder. + skips = [] + aux = x + for name, block in self.enc.items(): + if 'aux_down' in name: + aux = block(aux) + elif 'aux_skip' in name: + x = skips[-1] = x + block(aux) + elif 'aux_residual' in name: + x = skips[-1] = aux = (x + block(aux)) / np.sqrt(2) + else: + x = block(x, emb) if isinstance(block, UNetBlock) else block(x) + skips.append(x) + + # Decoder. + aux = None + tmp = None + for name, block in self.dec.items(): + if 'aux_up' in name: + aux = block(aux) + elif 'aux_norm' in name: + tmp = block(x) + elif 'aux_conv' in name: + tmp = block(silu(tmp)) + aux = tmp if aux is None else tmp + aux + else: + if x.shape[1] != block.in_channels: + x = torch.cat([x, skips.pop()], dim=1) + x = block(x, emb) + + if diffusion_cond is not None: + aux, _ = aux.chunk(2, dim=1) + # print(f'Output aux shape before reshape: {aux.shape}') + # print(f'Output aux shape after reshape: {aux.unsqueeze(2).shape}') + return aux.unsqueeze(2) # reshape from [B, C, H, W] to [B, C, 1, H, W] + + + + + + diff --git a/matey/models/vit.py b/matey/models/vit.py index 12b5095..ee82042 100644 --- a/matey/models/vit.py +++ b/matey/models/vit.py @@ -7,6 +7,8 @@ import numpy as np from einops import rearrange from .spacetime_modules import SpaceTimeBlock_all2all +from .time_modules import FourierEmbedding, PositionalEmbedding, Linear +from torch.nn.functional import silu from .basemodel import BaseModel from ..data_utils.shared_utils import normalize_spatiotemporal_persample, get_top_variance_patchids, normalize_spatiotemporal_persample_graph from ..utils import ForwardOptionsBase, TrainOptionsBase, densenodes_to_graphnodes @@ -37,6 +39,7 @@ def build_vit(params): replace_patch=getattr(params, 'replace_patch', True), hierarchical=getattr(params, 'hierarchical', None), use_linear=getattr(params, 'use_linear', False), + diffusion=getattr(params, 'diffusion', False) ) return model @@ -53,7 +56,7 @@ class ViT_all2all(BaseModel): sts_f """ def __init__(self, tokenizer_heads=None, embed_dim=768, 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", replace_patch=True, SR_ratio=[1,1,1], hierarchical=None, use_linear=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, use_linear=False, diffusion=False): 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, use_linear=use_linear) self.drop_path = drop_path @@ -71,6 +74,33 @@ def __init__(self, tokenizer_heads=None, embed_dim=768, num_heads=12, processor self.processor_blocks=processor_blocks self.replace_patch=replace_patch assert not (self.replace_patch and self.sts_model) + + self.diffusion=diffusion + if self.diffusion: + #from https://github.com/NVlabs/edm/blob/008a4e5316c8e3bfe61a62f874bddba254295afb/training/networks.py#L269 + #FIXME: @Paul, currently place holder pls check the proper setting of these variables + model_channels = 128 # Base multiplier for the number of channels. + channel_mult_emb = 4 # Multiplier for the dimensionality of the embedding vector. + embedding_type = 'positional' # Timestep embedding type: 'positional' for DDPM++, 'fourier' for NCSN++. + channel_mult_noise = 1 # Timestep embedding size: 1 for DDPM++, 2 for NCSN++. + emb_channels = model_channels * channel_mult_emb + noise_channels = model_channels * channel_mult_noise + init = dict(init_mode='xavier_uniform') + + self.map_noise = PositionalEmbedding(num_channels=noise_channels, endpoint=True) if embedding_type == 'positional' else FourierEmbedding(num_channels=noise_channels) + self.map_layer0 = Linear(in_features=noise_channels, out_features=emb_channels, **init) + # self.map_layer1 = Linear(in_features=emb_channels, out_features=embed_dim, **init) + self.map_layer1 = Linear(in_features=emb_channels, out_features=emb_channels, **init) + + # self.affine = nn.ModuleDict({}) + # for imod in range(self.nhlevels): + # self.affine[str(imod)] = Linear(in_features=emb_channels, out_features=embed_dim, **init) + self.affine = Linear(in_features=emb_channels, out_features=embed_dim, **init) + + # self.skip_projection = nn.ModuleDict({}) + # for imod in range(self.nhlevels): + # self.skip_projection[str(imod)] = Linear(in_features=embed_dim*2, out_features=embed_dim, **init) + self.skip_projection = Linear(in_features=embed_dim*2, out_features=embed_dim, **init) def expand_sts_model(self): """ Appends addition sts blocks""" @@ -131,11 +161,27 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase, train_opts: cond_input = opts.cond_input isgraph=opts.isgraph field_labels_out=opts.field_labels_out + sigma = getattr(opts, 'sigma', None) + diffusion_cond = getattr(opts, 'diffusion_cond', None) ################################################################## 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.map_noise(sigma) + # print(f'noise_labels shape: {sigma.shape}, emb shape: {emb.shape}') + emb = emb.reshape(emb.shape[0], 2, -1).flip(1).reshape(*emb.shape) # swap sin/cos + # print(f'after swap emb shape: {emb.shape}') + emb = silu(self.map_layer0(emb)) + # print(f'after first layer emb shape: {emb.shape}') + emb = silu(self.map_layer1(emb)) #in shape + # print(f'after second layer emb shape: {emb.shape}') + emb = self.affine(emb) + # print(f'after affine emb shape: {emb.shape}') + 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: x = data.x#[nnodes, T, C] @@ -174,6 +220,13 @@ 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, _, _, _, _, _, _, _ = self.get_patchsequence(diffusion_cond, state_labels, tkhead_name, refineind=refineind, blockdict=blockdict, ilevel=imod, isgraph=isgraph) + diffusion_cond = rearrange(diffusion_cond, 't b c ntoken_tot -> b c (t ntoken_tot)') + + # if self.diffusion: + # x_padding = x_padding + emb.unsqueeze(-1) + # Repeat the steps for conditioning if present if conditioning: assert self.sts_model == False @@ -194,11 +247,35 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase, train_opts: x_padding = x_padding + c if iblk==0: + if self.diffusion: + if diffusion_cond is not None: + # add diffusion_cond as additional tokens for the diffusion model to attend to + x_padding = torch.cat([x_padding, diffusion_cond], dim=2) + + # add noise embedding to input tokens as conditional information for diffusion model + x_padding = torch.cat([x_padding, emb.unsqueeze(-1)], dim=2) + + x_input = x_padding.clone() + x_padding = blk(x_padding, sequence_parallel_group=sequence_parallel_group, bcs=bcs, leadtime=leadtime, mask_padding=mask4attblk ) else: + if iblk==len(self.blocks)-1 and self.diffusion: + # for the last block, add skip connection from input tokens to facilitate training of diffusion model + x_padding = torch.cat([x_padding, x_input], dim=1) + local_batch = x_padding.shape[0] + x_padding = rearrange(x_padding, 'b c L -> (b L) c') + x_padding = self.skip_projection(x_padding) # Residual connection from input to the last block in the level + x_padding = rearrange(x_padding, '(b L) c -> b c L', b = local_batch) x_padding = blk(x_padding, sequence_parallel_group=sequence_parallel_group, bcs=bcs, leadtime=None, mask_padding=mask4attblk) #self.debug_nan(x_padding, message="attention block") ################################################################################ + + if self.diffusion: + x_padding = x_padding[:, :, :-1] # remove noise embedding from tokens after processing + if diffusion_cond is not None: + doubled_L = x_padding.shape[2] + x_padding = x_padding[:, :, :doubled_L//2] # remove diffusion_cond part from tokens after processing + x_padding = rearrange(x_padding, 'b c (t ntoken_tot) -> t b c ntoken_tot', t=T) ######## Decode ######## if self.sts_model: From 5b911336f33edceccbeddac4314cd9b3bccb4f91 Mon Sep 17 00:00:00 2001 From: Paul Laiu Date: Wed, 20 May 2026 10:47:27 -0400 Subject: [PATCH 02/12] added training and generation scripts for diffusion model --- examples/basic_generate.py | 48 +++ examples/submit_batch_diffusion.sh | 44 +++ examples/submit_batch_generate.sh | 56 ++++ matey/__init__.py | 2 +- matey/generate.py | 487 +++++++++++++++++++++++++++++ matey/train.py | 89 ++++-- matey/utils/forward_options.py | 4 +- matey/utils/training_utils.py | 34 +- 8 files changed, 735 insertions(+), 29 deletions(-) create mode 100644 examples/basic_generate.py create mode 100644 examples/submit_batch_diffusion.sh create mode 100644 examples/submit_batch_generate.sh create mode 100644 matey/generate.py 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/submit_batch_diffusion.sh b/examples/submit_batch_diffusion.sh new file mode 100644 index 0000000..fa0695e --- /dev/null +++ b/examples/submit_batch_diffusion.sh @@ -0,0 +1,44 @@ +#!/bin/bash +#SBATCH -A LRN037 +#SBATCH -J matey +#SBATCH -o %x-%j.out +#SBATCH -t 02: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_CIFAR10_diffusion_TT.yaml +# export yaml_config=./config/Demo_CIFAR10_diffusion_ViT.yaml +# export yaml_config=./config/Demo_CIFAR10_diffusion_UNet.yaml +# export yaml_config=./config/Demo_MW_diffusion_UNet.yaml +# export yaml_config=./config/Demo_MW_diffusion_TT.yaml +# export yaml_config=./config/Demo_MW_diffusion_ViT.yaml +export yaml_config=./config/Demo_MW_diffusion_TurbTMod.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..1da65a7 --- /dev/null +++ b/examples/submit_batch_generate.sh @@ -0,0 +1,56 @@ +#!/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_CIFAR10_3level_Ti_1e-3lr/basic_config/demo_diffusion/" +# export model_dir="./Demo_Diffusion_CIFAR10_UNet_1e-3lr/basic_config/demo_diffusion/" +# export model_dir="./Demo_Diffusion_CIFAR10_3level_Ti_newembskip_1e-3lr/basic_config/demo_diffusion/" + +# export model_dir="./Demo_Diffusion_MW_cond_S_newembskip_lt5/basic_config/demo_diffusion/" +# export model_dir="./Demo_Diffusion_MW_cond_S_newembskip_1e-3lr/basic_config/demo_diffusion/" + +# export model_dir="./Demo_Diffusion_MW_UNet_1e-3lr/basic_config/demo_diffusion/" +# export model_dir="./Demo_Diffusion_MW_cond_UNet_1e-3lr/basic_config/demo_diffusion/" +# export model_dir="./Demo_Diffusion_MW_cond_UNet_newemb_lt5/basic_config/demo_diffusion/" + +export model_dir="./Demo_Diffusion_MW_cond_TurbT_3level_S_lt5/basic_config/demo_diffusion/" + + +# export output_dir="./CIFAR10_generation_outputs/" +# export output_dir="./MW_generation_outputs/" +# export output_dir="./MW_cond_generation_outputs_batches/" +# export output_dir="./MW_cond_generation_outputs_batches_lt5/turbt/" +# export output_dir="./MW_cond_generation_outputs_batches_lt5/UNet/" +export output_dir="./MW_cond_generation_outputs_batches_lt5/TurbTMod/" + +##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..24ace59 --- /dev/null +++ b/matey/generate.py @@ -0,0 +1,487 @@ +# 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.turbt_modified import build_turbt_modified +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 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.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.cond_diffusion = getattr(self.params, "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")}') + + for sample_idx in range(num_samples): + + + init_inp = torch.randn(inp.shape).to(self.device) + + + self.model.eval() + + with torch.no_grad(): + + + sigma_min = 0.002 + sigma_max = 80 + num_steps = 18 + rho=7 + S_churn=0 + S_min=0 + S_max=float('inf') + S_noise=1 + + step_indices = torch.arange(num_steps, device=self.device) + + t_steps = (sigma_max ** (1 / rho) + step_indices / (num_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho + t_steps = torch.cat([t_steps, torch.zeros_like(t_steps[:1])]) # t_N = 0 + + # tar = tar.to(self.device) + # init_inp = rearrange(init_inp.to(self.device), 'b t c d h w -> t b c d h w') + + + x_next = init_inp * t_steps[0] + + for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])): # 0, ..., N-1 + x_cur = x_next + + # Increase noise temporarily. + gamma = min(S_churn / num_steps, np.sqrt(2) - 1) if S_min <= t_cur <= S_max else 0 + t_hat = t_cur + gamma * t_cur + x_hat = x_cur + (t_hat ** 2 - t_cur ** 2).sqrt() * S_noise * torch.randn_like(x_cur) + + + opts = ForwardOptionsBase( + imod=imod, + imod_bottom=imod_bottom , + tkhead_name=tkhead_name, + sequence_parallel_group=seq_group, + leadtime=leadtime, + blockdict=copy.deepcopy(blockdict), + cond_dict=copy.deepcopy(cond_dict), + cond_input=cond_input, + isgraph=isgraph, + field_labels_out= field_labels + ) + if self.cond_diffusion: + opts.diffusion_cond = rearrange(inp.to(self.device), 't b c d h w -> b t c d h w') # conditioning on input history + # Euler step. + # denoised = net(x_hat, t_hat, class_labels).to(torch.float64) + # denoised = net(x_hat, t_hat, None) + + # denoised = self.inference(x_hat, t_hat, field_labels, bcs, opts) + denoised = self.model(x_hat, t_hat.repeat(x_hat.shape[1]), field_labels, bcs, opts) + + # print(f"denoised shape: {denoised.shape}, x_hat shape: {x_hat.shape}, t_hat shape: {t_hat.shape}") + d_cur = (x_hat - denoised) / t_hat + x_next = x_hat + (t_next - t_hat) * d_cur + + # print(f"x_next shape: {x_next.shape}, d_cur shape: {d_cur.shape}, t_next shape: {t_next.shape}, t_hat shape: {t_hat.shape}") + # Apply 2nd order correction. + if i < num_steps - 1: + + opts = ForwardOptionsBase( + imod=imod, + imod_bottom=imod_bottom , + tkhead_name=tkhead_name, + sequence_parallel_group=seq_group, + leadtime=leadtime, + blockdict=copy.deepcopy(blockdict), + cond_dict=copy.deepcopy(cond_dict), + cond_input=cond_input, + isgraph=isgraph, + field_labels_out= field_labels + ) + if self.cond_diffusion: + opts.diffusion_cond = rearrange(inp.to(self.device), 't b c d h w -> b t c d h w') # conditioning on input history + # denoised = net(x_next, t_next, class_labels).to(torch.float64) + # denoised = net(x_next, t_next, None) + # denoised = self.inference(x_next, t_next, field_labels, bcs, opts) + denoised = self.model(x_next, t_next.repeat(x_next.shape[1]), field_labels, bcs, opts) + + d_prime = (x_next - denoised) / t_next + x_next = x_hat + (t_next - t_hat) * (0.5 * d_cur + 0.5 * d_prime) + + torch.save(x_next.cpu(), os.path.join(self.output_dir, f'generation_step_{i}_batch_{batch_idx}_sample{sample_idx}.pt')) + + output = x_next + + ###full resolution### + # residuals = output - tar + # torch.save(output.cpu(), os.path.join('Demo_Diffusion_CIFAR10_finemodel/basic_config/demo_diffusion/training_checkpoints/', 'generation_output.pt')) + # self.single_print(f'Generation output saved to {"Demo_Diffusion_CIFAR10_finemodel/basic_config/demo_diffusion/training_checkpoints/generation_output.pt"}') + torch.save(output.cpu(), os.path.join(self.output_dir, f'generation_output_batch_{batch_idx}_sample{sample_idx}.pt')) + self.single_print(f'Generation output saved to {os.path.join(self.output_dir, f"generation_output_batch_{batch_idx}_sample{sample_idx}.pt")}') + + + torch.cuda.empty_cache() + + + + def autoregressive_generate(self, seed=None, num_samples=1, num_steps=10): + if self.global_rank == 0: + summary(self.model) + self.single_print("Starting Generation Loop...") + + if seed is not None: + seed_value = seed + torch.manual_seed(seed_value) + + data_iter = iter(self.valid_data_loader) + for step_idx in range(num_steps): + + + if step_idx==0: + data = next(data_iter) + + if "graph" in data: + graphdata = data["graph"].to(self.device) + tar = graphdata.y #[nnodes, C_tar] + leadtime = graphdata.leadtime #[nnodes, 1] + 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 + if supportdata: + cond_input = data["cond_input"].to(self.device) + else: + cond_input = 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 + + 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")}') + + for sample_idx in range(num_samples): + + + init_inp = torch.randn(inp.shape).to(self.device) + + + self.model.eval() + + with torch.no_grad(): + + + sigma_min = 0.002 + sigma_max = 80 + num_steps = 18 + rho=7 + S_churn=0 + S_min=0 + S_max=float('inf') + S_noise=1 + + step_indices = torch.arange(num_steps, device=self.device) + + t_steps = (sigma_max ** (1 / rho) + step_indices / (num_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho + t_steps = torch.cat([t_steps, torch.zeros_like(t_steps[:1])]) # t_N = 0 + + # tar = tar.to(self.device) + # init_inp = rearrange(init_inp.to(self.device), 'b t c d h w -> t b c d h w') + + + x_next = init_inp * t_steps[0] + + for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])): # 0, ..., N-1 + x_cur = x_next + + # Increase noise temporarily. + gamma = min(S_churn / num_steps, np.sqrt(2) - 1) if S_min <= t_cur <= S_max else 0 + t_hat = t_cur + gamma * t_cur + x_hat = x_cur + (t_hat ** 2 - t_cur ** 2).sqrt() * S_noise * torch.randn_like(x_cur) + + + opts = ForwardOptionsBase( + imod=imod, + imod_bottom=imod_bottom , + tkhead_name=tkhead_name, + sequence_parallel_group=seq_group, + leadtime=leadtime, + blockdict=copy.deepcopy(blockdict), + cond_dict=copy.deepcopy(cond_dict), + cond_input=cond_input, + isgraph=isgraph, + field_labels_out= field_labels + ) + if self.cond_diffusion: + opts.diffusion_cond = rearrange(inp.to(self.device), 't b c d h w -> b t c d h w') # conditioning on input history + # Euler step. + # denoised = net(x_hat, t_hat, class_labels).to(torch.float64) + # denoised = net(x_hat, t_hat, None) + + # denoised = self.inference(x_hat, t_hat, field_labels, bcs, opts) + denoised = self.model(x_hat, t_hat.repeat(x_hat.shape[1]), field_labels, bcs, opts) + + # print(f"denoised shape: {denoised.shape}, x_hat shape: {x_hat.shape}, t_hat shape: {t_hat.shape}") + d_cur = (x_hat - denoised) / t_hat + x_next = x_hat + (t_next - t_hat) * d_cur + + # print(f"x_next shape: {x_next.shape}, d_cur shape: {d_cur.shape}, t_next shape: {t_next.shape}, t_hat shape: {t_hat.shape}") + # Apply 2nd order correction. + if i < num_steps - 1: + + opts = ForwardOptionsBase( + imod=imod, + imod_bottom=imod_bottom , + tkhead_name=tkhead_name, + sequence_parallel_group=seq_group, + leadtime=leadtime, + blockdict=copy.deepcopy(blockdict), + cond_dict=copy.deepcopy(cond_dict), + cond_input=cond_input, + isgraph=isgraph, + field_labels_out= field_labels + ) + if self.cond_diffusion: + opts.diffusion_cond = rearrange(inp.to(self.device), 't b c d h w -> b t c d h w') # conditioning on input history + # denoised = net(x_next, t_next, class_labels).to(torch.float64) + # denoised = net(x_next, t_next, None) + # denoised = self.inference(x_next, t_next, field_labels, bcs, opts) + denoised = self.model(x_next, t_next.repeat(x_next.shape[1]), field_labels, bcs, opts) + + d_prime = (x_next - denoised) / t_next + x_next = x_hat + (t_next - t_hat) * (0.5 * d_cur + 0.5 * d_prime) + + # torch.save(x_next.cpu(), os.path.join(self.output_dir, f'generation_output_step_{i}_batch_{batch_idx}_sample{sample_idx}.pt')) + + output = x_next + + ###full resolution### + # residuals = output - tar + # torch.save(output.cpu(), os.path.join('Demo_Diffusion_CIFAR10_finemodel/basic_config/demo_diffusion/training_checkpoints/', 'generation_output.pt')) + # self.single_print(f'Generation output saved to {"Demo_Diffusion_CIFAR10_finemodel/basic_config/demo_diffusion/training_checkpoints/generation_output.pt"}') + torch.save(output.cpu(), os.path.join(self.output_dir, f'generation_output_batch_{batch_idx}_sample{sample_idx}.pt')) + self.single_print(f'Generation output saved to {os.path.join(self.output_dir, f"generation_output_batch_{batch_idx}_sample{sample_idx}.pt")}') + + + torch.cuda.empty_cache() + + diff --git a/matey/train.py b/matey/train.py index 600d263..29eb7ef 100644 --- a/matey/train.py +++ b/matey/train.py @@ -21,6 +21,8 @@ from .models.svit import build_svit from .models.vit import build_vit from .models.turbt import build_turbt +from .models.turbt_modified import build_turbt_modified +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,7 +35,7 @@ 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: @@ -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.params.diffusion: + 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.params.diffusion: + 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) @@ -187,6 +194,8 @@ def initialize_model(self): self.model = build_vit(self.params).to(self.device) elif self.params.model_type == "turbt": self.model = build_turbt(self.params).to(self.device) + elif self.params.model_type == "turbt_modified": + self.model = build_turbt_modified(self.params).to(self.device) if self.params.compile: print('WARNING: BFLOAT NOT SUPPORTED IN SOME COMPILE OPS SO SWITCHING TO FLOAT16') @@ -529,7 +538,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 getattr(self.params, "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: @@ -544,10 +556,17 @@ def train_one_epoch(self): data_time = 0 data_start = self.timer.get_time() 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)} + if getattr(self.params, "diffusion", False): + logs = {'train_EDMloss': torch.zeros(1).to(self.device), + '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)} + else: + 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)} steps = 0 grad_logs = defaultdict(lambda: torch.zeros(1, device=self.device)) grad_counts = defaultdict(lambda: torch.zeros(1, device=self.device)) @@ -626,11 +645,26 @@ def train_one_epoch(self): field_labels_out= field_labels_out ) 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 getattr(self.params, "diffusion", False): + if getattr(self.params, "cond_diffusion", False): + opts.diffusion_cond = rearrange(inp.to(self.device), 't b c d h w -> b t c d h w') + 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) + else: + assert torch.all(inp[0] == tar), "For unconditioned diffusion model, input should be the target" + output, loss_field = self.model_forward(inp, 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 getattr(self.params, "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: @@ -706,10 +740,17 @@ def train_one_epoch(self): def validate_one_epoch(self, full=False, cutoff_skip=False): self.model.eval() self.single_print('STARTING VALIDATION!!!') - logs = {'valid_rmse': torch.zeros(1).to(self.device), - '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 getattr(self.params, "diffusion", False): + logs = {'valid_EDMloss': torch.zeros(1).to(self.device), + 'valid_rmse': torch.zeros(1).to(self.device), + 'valid_nrmse': torch.zeros(1).to(self.device), + 'valid_l1': torch.zeros(1).to(self.device), + 'valid_ssim': torch.zeros(1).to(self.device)} + else: + logs = {'valid_rmse': torch.zeros(1).to(self.device), + '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 cutoff_skip: return logs loss_dset_logs = {dataset.type: torch.zeros(1, device=self.device) for dataset in self.valid_dataset.sub_dsets} @@ -790,9 +831,19 @@ def validate_one_epoch(self, full=False, cutoff_skip=False): isgraph=isgraph, field_labels_out= field_labels_out ) - 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 getattr(self.params, "diffusion", False): + if getattr(self.params, "cond_diffusion", False): + opts.diffusion_cond = rearrange(inp.to(self.device), 't b c d h w -> b t c d h w') + 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) + else: + assert torch.all(inp[0] == tar), "For unconditioned diffusion model, input should be the target" + output, loss_field = self.model_forward(inp, 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 986eafc..4b93806 100644 --- a/matey/utils/forward_options.py +++ b/matey/utils/forward_options.py @@ -25,4 +25,6 @@ class ForwardOptionsBase: field_labels_out: Optional[Tensor] = None #adaptive tokenization (1 of 2 settings) refine_ratio: Optional[float] = None - imod_bottom: int = 0 #needed only by turbt \ No newline at end of file + imod_bottom: int = 0 #needed only by turbt + 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 3cdc60e..d35e6ef 100644 --- a/matey/utils/training_utils.py +++ b/matey/utils/training_utils.py @@ -11,6 +11,7 @@ from .visualization_utils import checking_data_pred_tar import copy + def preprocess_target(leadtime, ramping_warmup = False): """ #Inputs: @@ -59,13 +60,7 @@ def autoregressive_rollout(model, inp, field_labels, bcs, opts: ForwardOptionsBa # output: Model output after the final autoregressive step ([B, C, D, H, W]) # rollout_steps: Number of autoregressive steps performed. """ - is_constant = torch.all(opts.leadtime == opts.leadtime[0, 0]) - if is_constant: - rollout_steps = int(opts.leadtime[0,0].item()) - else: - rollout_steps = preprocess_target(opts.leadtime) - raise ValueError(f"Not expecting unequal leadtime across samples, {rollout_steps, opts.leadtime, is_constant}") - + rollout_steps = preprocess_target(opts.leadtime) x_t = inp if opts.isgraph: n_steps = x_t.x.shape[1] #[nnodes, T, C] @@ -158,7 +153,6 @@ def compute_loss_and_logs(output, tar, graphdata, logs, loss_logs, dset_type, pa logs['train_nrmse'] += log_nrmse loss_logs[dset_type] += loss.item() logs['train_rmse'] += residuals.pow(2).mean(spatial_dims).sqrt().mean() - return loss, log_nrmse def update_loss_logs_inplace_eval(output, tar, graphdata, logs, loss_dset_logs, loss_l1_dset_logs, loss_rmse_dset_logs, dset_type): @@ -191,3 +185,27 @@ 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 + +#---------------------------------------------------------------------------- From 8220ee5909c5bf51f114047e2040845325ab06a0 Mon Sep 17 00:00:00 2001 From: Paul Laiu Date: Wed, 20 May 2026 10:48:31 -0400 Subject: [PATCH 03/12] added config files for the four denoiser backbones --- examples/config/Demo_MW_diffusion_TT.yaml | 91 +++++++++++++++++++ .../config/Demo_MW_diffusion_TurbTMod.yaml | 71 +++++++++++++++ examples/config/Demo_MW_diffusion_UNet.yaml | 89 ++++++++++++++++++ examples/config/Demo_MW_diffusion_ViT.yaml | 88 ++++++++++++++++++ 4 files changed, 339 insertions(+) create mode 100644 examples/config/Demo_MW_diffusion_TT.yaml create mode 100644 examples/config/Demo_MW_diffusion_TurbTMod.yaml create mode 100644 examples/config/Demo_MW_diffusion_UNet.yaml create mode 100644 examples/config/Demo_MW_diffusion_ViT.yaml diff --git a/examples/config/Demo_MW_diffusion_TT.yaml b/examples/config/Demo_MW_diffusion_TT.yaml new file mode 100644 index 0000000..e5bf50f --- /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_S_newembskip_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: !!bool True + cond_diffusion: !!bool True + 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'], + # ['/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/config/Demo_MW_diffusion_TurbTMod.yaml b/examples/config/Demo_MW_diffusion_TurbTMod.yaml new file mode 100644 index 0000000..51d3733 --- /dev/null +++ b/examples/config/Demo_MW_diffusion_TurbTMod.yaml @@ -0,0 +1,71 @@ +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_TurbT_3level_S_lt5' # 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 + max_epochs: 200 + scheduler_epochs: -1 + epoch_size: 100 # Artificial epoch size + rescale_gradients: !!bool False # Activate hook that scales block gradients to norm 1 + optimizer: 'AdamW' + 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 # 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 + # NOTE: 'turbt_modified' must be registered in matey/models/diffusion_model.py + # (and train.py / generate.py / inference.py) alongside the existing 'turbt' entry. + model_type: 'turbt_modified' + diffusion: !!bool True + cond_diffusion: !!bool True + # Transformer size - S configuration: embed_dim/num_heads/processor_blocks = 384/6/12 + # Use 768/12/12 for B, or 1152/16/28 for XL. + embed_dim: 384 # Dimension of internal representation + num_heads: 6 # Number of attention heads + processor_blocks: 12 # Number of TurbTModified transformer blocks + tokenizer_heads: + - head_name: "tk-2D" + patch_size: [[1, 4, 4]] + hierarchical: + filtersize: 2 + nlevels: 3 + fixedupsample: !!bool False + linearupsample: !!bool False + tie_fields: !!bool False # Whether to use 1 embedding per field per data + 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: [ + ['/lustre/orion/lrn037/proj-shared/miniweather_thermals_tiny/train', 'thermalcollision2d', '', 'tk-2D'], + ] + valid_data_paths: [ + ['/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_UNet.yaml b/examples/config/Demo_MW_diffusion_UNet.yaml new file mode 100644 index 0000000..a837382 --- /dev/null +++ b/examples/config/Demo_MW_diffusion_UNet.yaml @@ -0,0 +1,89 @@ +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: 2 # 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_UNet_newemb_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: 'unet' + # 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: 192 # Dimension of internal representation - 192/384/768/1024 for Ti/S/B/L + num_heads: 3 # 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: !!bool True + cond_diffusion: !!bool True + 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/config/Demo_MW_diffusion_ViT.yaml b/examples/config/Demo_MW_diffusion_ViT.yaml new file mode 100644 index 0000000..75d3539 --- /dev/null +++ b/examples/config/Demo_MW_diffusion_ViT.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_ViT_S_newembskip_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: !!bool True + cond_diffusion: !!bool True + 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 + From 0e1b388b1f58eaaf968fec531e9ffc2db7e5b38b Mon Sep 17 00:00:00 2001 From: Paul Laiu Date: Wed, 20 May 2026 12:13:30 -0400 Subject: [PATCH 04/12] added new diffusion noise and condition embedding mechanism --- matey/models/turbt.py | 117 ++++++++++++++---------------------------- 1 file changed, 39 insertions(+), 78 deletions(-) diff --git a/matey/models/turbt.py b/matey/models/turbt.py index 3b57edc..a394e88 100644 --- a/matey/models/turbt.py +++ b/matey/models/turbt.py @@ -44,7 +44,7 @@ def build_turbt(params): replace_patch=getattr(params, 'replace_patch', True), hierarchical=getattr(params, 'hierarchical', None), notransposed=getattr(params, 'notransposed', False), - diffusion=getattr(params, 'diffusion', False) + diffusion=getattr(params, 'diffusion', False), ) return model @@ -119,34 +119,31 @@ 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)]) - - self.diffusion=diffusion + + self.diffusion = diffusion if self.diffusion: - #from https://github.com/NVlabs/edm/blob/008a4e5316c8e3bfe61a62f874bddba254295afb/training/networks.py#L269 - #FIXME: @Paul, currently place holder pls check the proper setting of these variables - model_channels = 128 # Base multiplier for the number of channels. - channel_mult_emb = 4 # Multiplier for the dimensionality of the embedding vector. - embedding_type = 'positional' # Timestep embedding type: 'positional' for DDPM++, 'fourier' for NCSN++. - channel_mult_noise = 1 # Timestep embedding size: 1 for DDPM++, 2 for NCSN++. + model_channels = 128 + channel_mult_emb = 4 + embedding_type = 'positional' + 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(num_channels=noise_channels, endpoint=True) if embedding_type == 'positional' else FourierEmbedding(num_channels=noise_channels) + + self.map_noise = (PositionalEmbedding(num_channels=noise_channels, endpoint=True) + if embedding_type == 'positional' + else FourierEmbedding(num_channels=noise_channels)) self.map_layer0 = Linear(in_features=noise_channels, out_features=emb_channels, **init) - # self.map_layer1 = Linear(in_features=emb_channels, out_features=embed_dim, **init) self.map_layer1 = Linear(in_features=emb_channels, out_features=emb_channels, **init) self.affine = nn.ModuleDict({}) for imod in range(self.nhlevels): self.affine[str(imod)] = Linear(in_features=emb_channels, out_features=embed_dim, **init) - # self.affine = Linear(in_features=emb_channels, out_features=embed_dim, **init) - self.skip_projection = nn.ModuleDict({}) + self.diffusion_cond_proj = nn.ModuleDict({}) for imod in range(self.nhlevels): - self.skip_projection[str(imod)] = Linear(in_features=embed_dim*2, out_features=embed_dim, **init) + self.diffusion_cond_proj[str(imod)] = Linear(in_features=embed_dim, out_features=embed_dim, **init) - 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}" @@ -254,32 +251,26 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase): field_labels_out=opts.field_labels_out sigma = getattr(opts, 'sigma', None) diffusion_cond = getattr(opts, 'diffusion_cond', None) - persample_normalize = False + 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.map_noise(sigma) - # print(f'noise_labels shape: {sigma.shape}, emb shape: {emb.shape}') - emb = emb.reshape(emb.shape[0], 2, -1).flip(1).reshape(*emb.shape) # swap sin/cos - # print(f'after swap emb shape: {emb.shape}') + emb = emb.reshape(emb.shape[0], 2, -1).flip(1).reshape(*emb.shape) # swap sin/cos emb = silu(self.map_layer0(emb)) - # print(f'after first layer emb shape: {emb.shape}') - emb = silu(self.map_layer1(emb)) #in shape - # print(f'after second layer emb shape: {emb.shape}') - emb = self.affine[str(imod)](emb) - # print(f'after affine emb shape: {emb.shape}') - if diffusion_cond is not None: - if 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 - # print(f"Paul debugging, imod = {imod}, diffusion_cond shape after rearrange: {diffusion_cond.shape}", flush=True) + emb = silu(self.map_layer1(emb)) + emb = self.affine[str(imod)](emb) # (B, embed_dim) + + 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: """ @@ -298,8 +289,7 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase): 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") 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") + #self.debug_nan(x, message="input after normalization") ################################################################################ if self.leadtime and leadtime is not None: leadtime = self.ltimeMLP[imod](leadtime) @@ -327,20 +316,23 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase): if self.cond_input and cond_input is not None: leadtime = self.inconMLP[imod](cond_input) if leadtime is None else leadtime+self.inconMLP[imod](cond_input) ########Encode and get patch sequences [B, C_emb, T*ntoken_len_tot]######## - # print(f"Paul debugging, diffusion_cond shape before get_patchsequence: {diffusion_cond.shape}, x shape: {x.shape}", flush=True) 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) - # print(f"Paul debugging, x shape after get_patchsequence: {x.shape}", flush=True) x = rearrange(x, 't b c ntoken_tot -> b c (t ntoken_tot)') - # print(f"Paul debugging, x shape after rearrange: {x.shape}", flush=True) + ################################################################################ if diffusion_cond is not None: - diffusion_cond, _, _, _, _, _, _, _ = self.get_patchsequence(diffusion_cond, state_labels, tkhead_name, refineind=refineind, blockdict=blockdict, ilevel=imod, isgraph=isgraph) - # print(f"Paul debugging, diffusion_cond shape after get_patchsequence: {diffusion_cond.shape}", flush=True) - diffusion_cond = rearrange(diffusion_cond, 't b c ntoken_tot -> b c (t ntoken_tot)') - # print(f"Paul debugging, diffusion_cond shape after get_patchsequence reshaping: {diffusion_cond.shape}", flush=True) - - # if self.diffusion: - # print("Paul debugging, sigma, emb, x", sigma.shape, emb.shape, x.shape, flush=True) - # x = x + emb.unsqueeze(-1) + 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 @@ -360,47 +352,16 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase): ps = self.tokenizer_ensemble_heads[imod][tkhead_name]["embed"][-1].patch_size nfact=4//ps[-1] """ - if diffusion_cond is not None: - # print(f"Paul debugging, before sequence_factor_short, diffusion_cond shape: {diffusion_cond.shape}, x shape: {x.shape}", flush=True) - diffusion_cond, _=self.sequence_factor_short(diffusion_cond, imod, tkhead_name, [T, D, H, W], nfact=nfact) - # print(f"Paul debugging, after sequence_factor_short, diffusion_cond shape: {diffusion_cond.shape}, x shape: {x.shape}", flush=True) - x, nfact=self.sequence_factor_short(x, imod, tkhead_name, [T, D, H, W], nfact=nfact) - for iblk, blk in enumerate(self.module_blocks[str(imod)]): if iblk==0: b_mod=x.shape[0] if not isgraph and leadtime is not None: leadtime = leadtime.repeat(b_mod // B, 1) - if self.diffusion: - if diffusion_cond is not None: - # add diffusion_cond as additional tokens for the diffusion model to attend to - x = torch.cat([x, diffusion_cond], dim=2) - - # add noise embedding to input tokens as conditional information for diffusion model - emb_mod = emb.repeat(b_mod // B, 1) - x = torch.cat([x, emb_mod.unsqueeze(-1)], dim=2) - - x_input = x.clone() x = blk(x, sequence_parallel_group=sequence_parallel_group, bcs=bcs, leadtime=leadtime, mask_padding=mask4attblk, local_att=local_att) else: - if iblk==len(self.module_blocks[str(imod)])-1 and self.diffusion: - # for the last block, add skip connection from input tokens to facilitate training of diffusion model - x = torch.cat([x, x_input], dim=1) - local_batch = x.shape[0] - x = rearrange(x, 'b c L -> (b L) c') - x = self.skip_projection[str(imod)](x) # Residual connection from input to the last block in the level - x = rearrange(x, '(b L) c -> b c L', b = local_batch) x = blk(x, sequence_parallel_group=sequence_parallel_group, bcs=bcs, leadtime=None, mask_padding=mask4attblk, local_att=local_att) - #self.debug_nan(x_padding, message="attention block") - - if self.diffusion: - x = x[:, :, :-1] # remove noise embedding from tokens after processing - if diffusion_cond is not None: - doubled_L = x.shape[2] - x = x[:, :, :doubled_L//2] # remove diffusion_cond part from tokens after processing - if local_att: #nfact=4//ps[-1] x=self.sequence_factor_long(x, imod, tkhead_name, [T, D, H, W], nfact=nfact) From c5371ecbc6fb1e2e073a49996057dbeb0bdcb72d Mon Sep 17 00:00:00 2001 From: Paul Laiu Date: Wed, 20 May 2026 12:30:23 -0400 Subject: [PATCH 05/12] removed redundant options; before refactoring --- examples/config/Demo_MW_diffusion_TT.yaml | 9 +- .../config/Demo_MW_diffusion_TurbTMod.yaml | 71 --- examples/config/Demo_MW_diffusion_UNet.yaml | 89 ---- examples/submit_batch_diffusion.sh | 9 +- examples/submit_batch_generate.sh | 16 +- matey/generate.py | 12 +- matey/models/__init__.py | 3 +- matey/models/diffusion_model.py | 8 - matey/models/turbt_modified.py | 454 ------------------ matey/models/unet.py | 354 -------------- matey/models/vit.py | 79 +-- matey/train.py | 3 - 12 files changed, 9 insertions(+), 1098 deletions(-) delete mode 100644 examples/config/Demo_MW_diffusion_TurbTMod.yaml delete mode 100644 examples/config/Demo_MW_diffusion_UNet.yaml delete mode 100644 matey/models/turbt_modified.py delete mode 100644 matey/models/unet.py diff --git a/examples/config/Demo_MW_diffusion_TT.yaml b/examples/config/Demo_MW_diffusion_TT.yaml index e5bf50f..015adad 100644 --- a/examples/config/Demo_MW_diffusion_TT.yaml +++ b/examples/config/Demo_MW_diffusion_TT.yaml @@ -8,7 +8,7 @@ basic_config: &basic_config 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_S_newembskip_lt5' # Output path + exp_dir: 'Demo_Diffusion_MW_cond_TurbT_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 @@ -71,10 +71,7 @@ basic_config: &basic_config #['/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'], - + ['/lustre/orion/lrn037/proj-shared/miniweather_thermals_tiny/train', 'thermalcollision2d', '', 'tk-2D'], ] valid_data_paths: [ #['/home/6pz/data/PDEBench_Data/2D/shallow-water', 'swe', ''], @@ -84,8 +81,6 @@ basic_config: &basic_config #['/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/config/Demo_MW_diffusion_TurbTMod.yaml b/examples/config/Demo_MW_diffusion_TurbTMod.yaml deleted file mode 100644 index 51d3733..0000000 --- a/examples/config/Demo_MW_diffusion_TurbTMod.yaml +++ /dev/null @@ -1,71 +0,0 @@ -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_TurbT_3level_S_lt5' # 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 - max_epochs: 200 - scheduler_epochs: -1 - epoch_size: 100 # Artificial epoch size - rescale_gradients: !!bool False # Activate hook that scales block gradients to norm 1 - optimizer: 'AdamW' - 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 # 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 - # NOTE: 'turbt_modified' must be registered in matey/models/diffusion_model.py - # (and train.py / generate.py / inference.py) alongside the existing 'turbt' entry. - model_type: 'turbt_modified' - diffusion: !!bool True - cond_diffusion: !!bool True - # Transformer size - S configuration: embed_dim/num_heads/processor_blocks = 384/6/12 - # Use 768/12/12 for B, or 1152/16/28 for XL. - embed_dim: 384 # Dimension of internal representation - num_heads: 6 # Number of attention heads - processor_blocks: 12 # Number of TurbTModified transformer blocks - tokenizer_heads: - - head_name: "tk-2D" - patch_size: [[1, 4, 4]] - hierarchical: - filtersize: 2 - nlevels: 3 - fixedupsample: !!bool False - linearupsample: !!bool False - tie_fields: !!bool False # Whether to use 1 embedding per field per data - 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: [ - ['/lustre/orion/lrn037/proj-shared/miniweather_thermals_tiny/train', 'thermalcollision2d', '', 'tk-2D'], - ] - valid_data_paths: [ - ['/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_UNet.yaml b/examples/config/Demo_MW_diffusion_UNet.yaml deleted file mode 100644 index a837382..0000000 --- a/examples/config/Demo_MW_diffusion_UNet.yaml +++ /dev/null @@ -1,89 +0,0 @@ -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: 2 # 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_UNet_newemb_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: 'unet' - # 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: 192 # Dimension of internal representation - 192/384/768/1024 for Ti/S/B/L - num_heads: 3 # 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: !!bool True - cond_diffusion: !!bool True - 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 index fa0695e..8d2b037 100644 --- a/examples/submit_batch_diffusion.sh +++ b/examples/submit_batch_diffusion.sh @@ -13,13 +13,10 @@ export OMP_NUM_THREADS=1 export master_node=$SLURMD_NODENAME export config="basic_config" export run_name="demo_diffusion" -# export yaml_config=./config/Demo_CIFAR10_diffusion_TT.yaml -# export yaml_config=./config/Demo_CIFAR10_diffusion_ViT.yaml -# export yaml_config=./config/Demo_CIFAR10_diffusion_UNet.yaml -# export yaml_config=./config/Demo_MW_diffusion_UNet.yaml -# export yaml_config=./config/Demo_MW_diffusion_TT.yaml + +export yaml_config=./config/Demo_MW_diffusion_TT.yaml # export yaml_config=./config/Demo_MW_diffusion_ViT.yaml -export yaml_config=./config/Demo_MW_diffusion_TurbTMod.yaml + ##conda env with rocm 6.0.0 #module load rocm/6.0.0 diff --git a/examples/submit_batch_generate.sh b/examples/submit_batch_generate.sh index 1da65a7..6a329a9 100644 --- a/examples/submit_batch_generate.sh +++ b/examples/submit_batch_generate.sh @@ -12,26 +12,12 @@ export OMP_NUM_THREADS=1 export master_node=$SLURMD_NODENAME export config="basic_config" -# export model_dir="./Demo_Diffusion_CIFAR10_3level_Ti_1e-3lr/basic_config/demo_diffusion/" -# export model_dir="./Demo_Diffusion_CIFAR10_UNet_1e-3lr/basic_config/demo_diffusion/" -# export model_dir="./Demo_Diffusion_CIFAR10_3level_Ti_newembskip_1e-3lr/basic_config/demo_diffusion/" -# export model_dir="./Demo_Diffusion_MW_cond_S_newembskip_lt5/basic_config/demo_diffusion/" -# export model_dir="./Demo_Diffusion_MW_cond_S_newembskip_1e-3lr/basic_config/demo_diffusion/" - -# export model_dir="./Demo_Diffusion_MW_UNet_1e-3lr/basic_config/demo_diffusion/" -# export model_dir="./Demo_Diffusion_MW_cond_UNet_1e-3lr/basic_config/demo_diffusion/" -# export model_dir="./Demo_Diffusion_MW_cond_UNet_newemb_lt5/basic_config/demo_diffusion/" export model_dir="./Demo_Diffusion_MW_cond_TurbT_3level_S_lt5/basic_config/demo_diffusion/" +export output_dir="./MW_cond_generation_outputs_batches_lt5/turbt/" -# export output_dir="./CIFAR10_generation_outputs/" -# export output_dir="./MW_generation_outputs/" -# export output_dir="./MW_cond_generation_outputs_batches/" -# export output_dir="./MW_cond_generation_outputs_batches_lt5/turbt/" -# export output_dir="./MW_cond_generation_outputs_batches_lt5/UNet/" -export output_dir="./MW_cond_generation_outputs_batches_lt5/TurbTMod/" ##conda env with rocm 6.0.0 #module load rocm/6.0.0 diff --git a/matey/generate.py b/matey/generate.py index 24ace59..ccd1411 100644 --- a/matey/generate.py +++ b/matey/generate.py @@ -12,7 +12,6 @@ from .models.svit import build_svit from .models.vit import build_vit from .models.turbt import build_turbt -from .models.turbt_modified import build_turbt_modified 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 @@ -100,16 +99,7 @@ def initialize_model(self): if self.params.diffusion: self.model = build_diffusion_model(self.params).to(self.device) else: - if 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) - elif self.params.model_type == "vit_all2all": - self.model = build_vit(self.params).to(self.device) - elif self.params.model_type == "turbt": - self.model = build_turbt(self.params).to(self.device) - elif self.params.model_type == "turbt_modified": - self.model = build_turbt_modified(self.params).to(self.device) + raise NotImplementedError("Only diffusion model generation is implemented currently. Please set params.diffusion to True.") if dist.is_initialized() and self.params.use_ddp: diff --git a/matey/models/__init__.py b/matey/models/__init__.py index eae28a6..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 -from .unet import build_unet, UNet -__all__ = ["build_avit", "build_svit", "build_vit","build_turbt", "AViT","sViT_all2all","ViT_all2all","TurbT", "build_unet", "UNet"] +__all__ = ["build_avit", "build_svit", "build_vit","build_turbt", "AViT","sViT_all2all","ViT_all2all","TurbT"] diff --git a/matey/models/diffusion_model.py b/matey/models/diffusion_model.py index cd812d4..3f8ef66 100644 --- a/matey/models/diffusion_model.py +++ b/matey/models/diffusion_model.py @@ -6,8 +6,6 @@ from .svit import build_svit from .vit import build_vit from .turbt import build_turbt -from .turbt_modified import build_turbt_modified -from .unet import build_unet def build_diffusion_model(params): model = EDMPrecond(params=params) @@ -43,12 +41,6 @@ def __init__(self, elif params.model_type == 'turbt': self.model = build_turbt(params) self.tokenizer_heads_params = self.model.tokenizer_heads_params - elif params.model_type == 'turbt_modified': - self.model = build_turbt_modified(params) - self.tokenizer_heads_params = self.model.tokenizer_heads_params - elif params.model_type == 'unet': - self.model = build_unet(params) - self.tokenizer_heads_params = {'CIFAR10': [[1, 1, 1]]} else: raise ValueError(f"Unknown diffusion model type: {params.model_type}") diff --git a/matey/models/turbt_modified.py b/matey/models/turbt_modified.py deleted file mode 100644 index f2eae1b..0000000 --- a/matey/models/turbt_modified.py +++ /dev/null @@ -1,454 +0,0 @@ -# SPDX-License-Identifier: MIT -# SPDX-FileCopyrightText: 2026 UT-Battelle, LLC -# This file is part of the MATEY Project. - -""" -TurbT modified to eliminate patch artifacts in diffusion mode. - -Three changes relative to turbt.py, all inspired by DiffiT's design: - -1. Per-block broadcast noise conditioning (replaces single appended time token) - - Original: noise embedding appended as one extra token at iblk==0; attention - must distribute it across all spatial tokens, giving non-uniform conditioning - that varies with distance from the token boundary → visible patch artifacts. - - Modified: emb_mod is broadcast-added to ALL tokens via the existing leadtime - mechanism inside AttentionBlock_all2all at EVERY block, giving spatially - uniform noise conditioning (same as DiffiT's per-token QKV bias approach). - -2. Conditioning token fusion via learned projection + addition (replaces appended - diffusion_cond tokens that doubled the sequence length) - - Original: diffusion_cond tokens concatenated to the sequence, creating a - seam at the midpoint; tokens near the boundary see a different context than - tokens far from it → boundary artifacts aligned with patch grid. - - Modified: diffusion_cond tokens are projected to embed_dim and added - element-wise to the main token stream before the block loop (analogous to - DiffiT/DiffiTMATEY's channel-concatenation approach, applied in token space). - -All non-diffusion code paths (weather forecasting, graph, hierarchical) are -unchanged. -""" - -import torch -import torch.nn as nn -import torch.nn.functional as F -import numpy as np -from einops import rearrange -from .spacetime_modules import SpaceTimeBlock_all2all -from .basemodel import BaseModel -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 -from .time_modules import FourierEmbedding, PositionalEmbedding, Linear -from torch.nn.functional import silu -import sys, copy -from operator import mul -from functools import reduce -from ..utils.forward_options import ForwardOptionsBase -import torch.distributed as dist -from ..utils import densenodes_to_graphnodes - - -def build_turbt_modified(params): - """ Builds modified model from parameter file. Same signature as build_turbt. """ - model = TurbTModified( - tokenizer_heads=params.tokenizer_heads, - embed_dim=params.embed_dim, - num_heads=params.num_heads, - processor_blocks=params.processor_blocks, - n_states=params.n_states, - sts_model=getattr(params, 'sts_model', False), - sts_train=getattr(params, 'sts_train', False), - leadtime=hasattr(params, "leadtime_max") and params.leadtime_max >= 0, - cond_input=getattr(params, 'supportdata', False), - n_steps=params.n_steps, - bias_type=params.bias_type, - replace_patch=getattr(params, 'replace_patch', True), - hierarchical=getattr(params, 'hierarchical', None), - notransposed=getattr(params, 'notransposed', False), - diffusion=getattr(params, 'diffusion', False), - ) - return model - - -class TurbTModified(BaseModel): - """ - TurbT with DiffiT-inspired diffusion conditioning to eliminate patch artifacts. - - Identical to TurbT for all non-diffusion code paths. - """ - - def __init__(self, tokenizer_heads=None, embed_dim=768, num_heads=12, - processor_blocks=8, n_states=6, 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=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, - ) - self.drop_path = drop_path - self.dp = np.linspace(0, drop_path, processor_blocks) - self.module_blocks = nn.ModuleDict({}) - self.sts_model = sts_model - self.sts_train = sts_train - - self.num_heads = num_heads - self.n_steps = n_steps - self.processor_blocks = processor_blocks - self.replace_patch = replace_patch - assert not (self.replace_patch and self.sts_model) - - self.upscale_factors = [1] - self.module_upscale = nn.ModuleDict({}) - self.module_upscale_space = nn.ModuleDict({}) - self.module_upscale_space2D = nn.ModuleDict({}) - - self.hierarchical = False - self.datafilter_kernel = None - if hierarchical is not None: - self.hierarchical = True - filtersize = hierarchical["filtersize"] - self.datafilter_kernel = construct_filterkernel(filtersize) - self.datafilter_kernel2D = construct_filterkernel2D(filtersize) - self.filtersize = filtersize - self.nhlevels = hierarchical["nlevels"] - self.upscale_factors = [1] + [self.filtersize for _ in range(self.nhlevels - 1)] - - for imod, upscalefactor in enumerate(self.upscale_factors): - if hierarchical["fixedupsample"]: - self.module_upscale_space[str(imod)] = nn.Upsample( - scale_factor=(upscalefactor, upscalefactor, upscalefactor), - mode='trilinear', align_corners=True) - self.module_upscale_space2D[str(imod)] = nn.Upsample( - scale_factor=(1, upscalefactor, upscalefactor), - mode='trilinear', align_corners=True) - elif hierarchical["linearupsample"]: - self.module_upscale_space[str(imod)] = torch.nn.Sequential( - nn.Upsample(scale_factor=(upscalefactor, upscalefactor, upscalefactor), - mode='trilinear', align_corners=True), - nn.Conv3d(n_states, n_states, - kernel_size=(upscalefactor, upscalefactor, upscalefactor), - stride=1, padding="same", bias=True, padding_mode="reflect"), - nn.InstanceNorm3d(n_states, affine=True)) - self.module_upscale_space2D[str(imod)] = torch.nn.Sequential( - nn.Upsample(scale_factor=(1, upscalefactor, upscalefactor), - mode='trilinear', align_corners=True), - nn.Conv3d(n_states, n_states, - kernel_size=(1, upscalefactor, upscalefactor), - stride=1, padding="same", bias=True, padding_mode="reflect"), - nn.InstanceNorm3d(n_states, affine=True)) - else: - self.module_upscale_space[str(imod)] = UpsampleinSpace( - patch_size=[upscalefactor, upscalefactor, upscalefactor], channels=n_states) - self.module_upscale_space2D[str(imod)] = UpsampleinSpace( - patch_size=[1, upscalefactor, upscalefactor], channels=n_states) - - if imod == 0: - 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)]) - else: - 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)]) - - self.diffusion = diffusion - if self.diffusion: - model_channels = 128 - channel_mult_emb = 4 - embedding_type = 'positional' - 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(num_channels=noise_channels, endpoint=True) - if embedding_type == 'positional' - else FourierEmbedding(num_channels=noise_channels)) - self.map_layer0 = Linear(in_features=noise_channels, out_features=emb_channels, **init) - self.map_layer1 = Linear(in_features=emb_channels, out_features=emb_channels, **init) - - self.affine = nn.ModuleDict({}) - for imod in range(self.nhlevels): - self.affine[str(imod)] = Linear(in_features=emb_channels, out_features=embed_dim, **init) - - # NEW: per-level learned projection for fusing diffusion_cond tokens into - # the main token stream via element-wise addition (replaces token appending). - self.diffusion_cond_proj = nn.ModuleDict({}) - for imod in range(self.nhlevels): - self.diffusion_cond_proj[str(imod)] = Linear(in_features=embed_dim, out_features=embed_dim, **init) - - # ------------------------------------------------------------------ - # Helpers copied verbatim from TurbT - # ------------------------------------------------------------------ - - def filterdata(self, data, blockdict=None): - assert data.ndim == 6, f"unkown tensor shape in filter_data, {data.shape}" - with torch.no_grad(): - kernel_size = self.filtersize - T, B, C, D, H, W = data.shape - data = rearrange(data, 't b c d h w -> (t b c) d h w') - if D == 1: - kernel = self.datafilter_kernel2D - filtered = F.conv3d(data[:, None, :, :, :], kernel.to(data.device), - stride=(1, kernel_size, kernel_size)) - else: - kernel = self.datafilter_kernel - filtered = F.conv3d(data[:, None, :, :, :], kernel.to(data.device), - stride=kernel_size) - filtered = rearrange(filtered, '(t b c) c1 d h w -> t b (c c1) d h w', t=T, b=B, c=C) - if blockdict is not None: - assert [D, H, W] == blockdict["Ind_dim"], f"(D,H,W),{(D,H,W)}, {blockdict['Ind_dim']}" - if D == 1: - blockdict["Ind_dim"] = [D, H // kernel_size, W // kernel_size] - else: - blockdict["Ind_dim"] = [D // kernel_size, H // kernel_size, W // kernel_size] - return filtered, blockdict - - def upsampeldata(self, data, imod): - B, C, D, H, W = data.shape - if D == 1: - data_upsample = self.module_upscale_space2D[str(imod)](data) - else: - data_upsample = self.module_upscale_space[str(imod)](data) - return data_upsample - - def sequence_factor_short(self, x, ilevel, tkhead_name, tspace_dims, nfact=2): - B, C, TL = x.shape - embed_ensemble = self.tokenizer_ensemble_heads[ilevel][tkhead_name]["embed"] - ntokendim = [] - ps_c = embed_ensemble[-1].patch_size - for idim, dim in enumerate(tspace_dims[1:]): - ntokendim.append(dim // ps_c[idim]) - assert TL == tspace_dims[0] * reduce(mul, ntokendim), \ - f"{TL}, {tspace_dims}, {ntokendim}" - d, h, w = ntokendim - if h // nfact < 4: - nfact = max(1, h // 4) - if nfact < 2: - return x, nfact - nfactd = 1 if d == 1 else nfact - x = rearrange(x, 'b c (t d h w) -> b c t d h w', - t=tspace_dims[0], d=d, h=h, w=w) - x = x.unfold(3, d // nfactd, d // nfactd).unfold(4, h // nfact, h // nfact).unfold(5, w // nfact, w // nfact) - x = rearrange(x, 'b c t nd nh nw d h w -> (b nd nh nw) c (t d h w)') - return x, nfact - - def sequence_factor_long(self, x, ilevel, tkhead_name, tspace_dims, nfact=2): - if nfact < 2: - return x - B, C, TL = x.shape - embed_ensemble = self.tokenizer_ensemble_heads[ilevel][tkhead_name]["embed"] - ntokendim = [] - ps_c = embed_ensemble[-1].patch_size - for idim, dim in enumerate(tspace_dims[1:]): - ntokendim.append(dim // ps_c[idim]) - d, h, w = ntokendim - nfactd = 1 if d == 1 else nfact - assert TL * (nfactd * nfact * nfact) == tspace_dims[0] * reduce(mul, ntokendim), \ - f"{TL}, {tspace_dims}, {ntokendim}, {nfact, nfactd}" - x = rearrange(x, '(b nd nh nw) c (t d h w) -> b c t nd nh nw d h w', - b=B // (nfactd * nfact * nfact), nd=nfactd, nh=nfact, nw=nfact, - d=d // nfactd, h=h // nfact, w=w // nfact) - x = rearrange(x, 'b c t nd nh nw d h w -> b c t (nd d) (nh h) (nw w)') - x = rearrange(x, 'b c t d h w -> b c (t d h w)') - return x - - # ------------------------------------------------------------------ - # Forward - # ------------------------------------------------------------------ - - def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase): - # ---- unpack arguments ------------------------------------------------ - imod = opts.imod - imod_bottom = opts.imod_bottom - tkhead_name = opts.tkhead_name - sequence_parallel_group = opts.sequence_parallel_group - leadtime = opts.leadtime - blockdict = opts.blockdict - refine_ratio = opts.refine_ratio - cond_input = opts.cond_input - isgraph = opts.isgraph - field_labels_out = opts.field_labels_out - sigma = getattr(opts, 'sigma', None) - diffusion_cond = getattr(opts, 'diffusion_cond', None) - persample_normalize = False - # ---- argument checks ------------------------------------------------- - if refine_ratio is not None: - raise ValueError("Adaptive tokenization is not set up/tested yet in TurbTModified") - - if field_labels_out is None: - field_labels_out = state_labels - - # ---- noise embedding (diffusion only) -------------------------------- - if self.diffusion: - emb = self.map_noise(sigma) - emb = emb.reshape(emb.shape[0], 2, -1).flip(1).reshape(*emb.shape) # swap sin/cos - emb = silu(self.map_layer0(emb)) - emb = silu(self.map_layer1(emb)) - emb = self.affine[str(imod)](emb) # (B, embed_dim) - - # Rearrange diffusion_cond to (T, B, C, D, H, W) at finest level. - # At coarser levels it arrives pre-filtered in the right format. - 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 - - # ---- graph path (unchanged from TurbT) -------------------------------- - if isgraph: - x = data.x - edge_index = data.edge_index - batch = data.batch - T = x.shape[1] - x, data_mean, data_std = normalize_spatiotemporal_persample_graph(x, batch) - refineind = None - x = (x, batch, edge_index) - else: - x = data - - if imod < self.nhlevels - 1: - # Filter x and diffusion_cond together so they share the same - # coarse grid (unchanged from TurbT). - if diffusion_cond is not None: - concat_data = torch.cat([x, diffusion_cond], dim=2) - concat_data, blockdict = self.filterdata(concat_data, blockdict=blockdict) - x = concat_data[:, :, :x.shape[2], :, :, :] - diffusion_cond = concat_data[:, :, x.shape[2]:, :, :, :] - opts.diffusion_cond = diffusion_cond - opts.blockdict = blockdict - else: - x, blockdict = self.filterdata(x, blockdict=blockdict) - opts.blockdict = blockdict - - if imod > imod_bottom: - opts.imod -= 1 - x_pred = self.forward(x, state_labels, bcs, opts) - - T, B, _, D, H, W = x.shape - if persample_normalize: - x, data_mean, data_std = normalize_spatiotemporal_persample(x) - - # ---- leadtime / conditional input (unchanged) ------------------------- - if self.leadtime and leadtime is not None: - leadtime = self.ltimeMLP[imod](leadtime) - else: - leadtime = None - if self.cond_input and cond_input is not None: - leadtime = (self.inconMLP[imod](cond_input) - if leadtime is None - else leadtime + self.inconMLP[imod](cond_input)) - - # ---- tokenize x ------------------------------------------------------- - x, patch_ids, patch_ids_ref, mask_padding, _, _, tposarea_padding, _ = \ - self.get_patchsequence(x, state_labels, tkhead_name, refineind=None, - blockdict=blockdict, ilevel=imod, isgraph=isgraph) - x = rearrange(x, 't b c ntoken_tot -> b c (t ntoken_tot)') - - # ---- fuse noise embedding and diffusion_cond into x (CHANGED) -------- - # Both the noise embedding (emb) and diffusion_cond are combined and - # added into x before the block loop. When diffusion_cond is present, - # emb is added to the cond tokens before projection so both signals are - # fused in one learned step. When diffusion_cond is absent, emb is - # broadcast directly into x. - 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) - - # ---- positional bias (unchanged) -------------------------------------- - 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) - posbias = rearrange(posbias, 'b t L c -> b c (t L)') - x = x + posbias - del posbias - - # ---- attention mask --------------------------------------------------- - mask4attblk = None if (mask_padding is not None and mask_padding.all()) else mask_padding - - # ---- local attention windowing (CHANGED: no separate cond handling) --- - # diffusion_cond has already been fused into x, so only x needs windowing. - local_att = not isgraph and imod > imod_bottom - if local_att: - nfact = (max(2 ** (2 * (imod - imod_bottom)) // blockdict["nproc_blocks"][-1], 1) - if blockdict is not None - else max(2 ** (2 * (imod - imod_bottom)), 1)) - x, nfact = self.sequence_factor_short(x, imod, tkhead_name, [T, D, H, W], nfact=nfact) - - # ---- transformer block loop (CHANGED) --------------------------------- - # Key changes: - # (a) No time embedding token appended to the sequence. - # (b) Noise embedding is fused into x before the loop (not via leadtime). - # (c) leadtime passed only at iblk==0, None thereafter (original behaviour). - for iblk, blk in enumerate(self.module_blocks[str(imod)]): - if iblk == 0: - b_mod = x.shape[0] - if not isgraph and leadtime is not None: - leadtime = leadtime.repeat(b_mod // B, 1) - - block_cond = leadtime if iblk == 0 else None - - x = blk(x, sequence_parallel_group=sequence_parallel_group, - bcs=bcs, leadtime=block_cond, mask_padding=mask4attblk, - local_att=local_att) - - # No time/cond tokens to strip — they were never appended. - - # ---- un-window (unchanged) -------------------------------------------- - if local_att: - x = self.sequence_factor_long(x, imod, tkhead_name, [T, D, H, W], nfact=nfact) - - # ---- decode ----------------------------------------------------------- - x = rearrange(x, 'b c (t ntoken_tot) -> t b c ntoken_tot', t=T) - - if isgraph: - x = rearrange(x, 't b c ntoken_tot -> b ntoken_tot t c') - x = densenodes_to_graphnodes(x, mask_padding) - x = (x, batch, edge_index) - D, H, W = -1, -1, -1 - - x = self.get_spatiotemporalfromsequence( - x, patch_ids, patch_ids_ref, [D, H, W], tkhead_name, - ilevel=imod, isgraph=isgraph) - - if isgraph: - node_ft, batch, edge_index = x - x = node_ft[:, :, field_labels_out[0]] - N = x.shape[0] - mask = torch.isin(state_labels[0], field_labels_out[0]) - mean_node = data_mean[batch].view(N, 1, -1)[:, :, mask] - std_node = data_std[batch].view(N, 1, -1)[:, :, mask] - x = x * std_node + mean_node - return x[:, -1, :] - - # ---- hierarchical upsampling (unchanged) ------------------------------ - x_correct = x[-1] - del x - if imod > imod_bottom: - x_filter = self.filterdata(x_correct[None, ...])[0][-1] - filtered_eps = self.upsampeldata(x_filter, imod) - x_correct = x_correct - filtered_eps - x_pred = self.upsampeldata(x_pred, imod) - x_correct = x_correct + x_pred - - if imod == self.nhlevels - 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], ...] - - return x_correct diff --git a/matey/models/unet.py b/matey/models/unet.py deleted file mode 100644 index c2659f5..0000000 --- a/matey/models/unet.py +++ /dev/null @@ -1,354 +0,0 @@ -# SPDX-License-Identifier: MIT -# SPDX-FileCopyrightText: 2026 UT-Battelle, LLC -# This file is part of the MATEY Project. - -import torch -import torch.nn as nn -import torch.nn.functional as F -import numpy as np -from einops import rearrange -from .spacetime_modules import SpaceTimeBlock_all2all -from .basemodel import BaseModel -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 -from .time_modules import FourierEmbedding, PositionalEmbedding, Linear -from torch.nn.functional import silu -import sys, copy -from operator import mul -from functools import reduce -from ..utils.forward_options import ForwardOptionsBase -import torch.distributed as dist -from ..utils import densenodes_to_graphnodes - - -def build_unet(params): - """ Builds model from parameter file. - """ - model = UNet() - return model - -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(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 - - -class Conv2d(torch.nn.Module): - def __init__(self, - in_channels, out_channels, kernel, bias=True, up=False, down=False, - resample_filter=[1,1], fused_resample=False, init_mode='kaiming_normal', init_weight=1, init_bias=0, - ): - assert not (up and down) - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.up = up - self.down = down - self.fused_resample = fused_resample - init_kwargs = dict(mode=init_mode, fan_in=in_channels*kernel*kernel, fan_out=out_channels*kernel*kernel) - self.weight = torch.nn.Parameter(weight_init([out_channels, in_channels, kernel, kernel], **init_kwargs) * init_weight) if kernel else None - self.bias = torch.nn.Parameter(weight_init([out_channels], **init_kwargs) * init_bias) if kernel and bias else None - f = torch.as_tensor(resample_filter, dtype=torch.float32) - f = f.ger(f).unsqueeze(0).unsqueeze(1) / f.sum().square() - self.register_buffer('resample_filter', f if up or down else None) - - def forward(self, x): - w = self.weight.to(x.dtype) if self.weight is not None else None - b = self.bias.to(x.dtype) if self.bias is not None else None - f = self.resample_filter.to(x.dtype) if self.resample_filter is not None else None - w_pad = w.shape[-1] // 2 if w is not None else 0 - f_pad = (f.shape[-1] - 1) // 2 if f is not None else 0 - - if self.fused_resample and self.up and w is not None: - x = torch.nn.functional.conv_transpose2d(x, f.mul(4).tile([self.in_channels, 1, 1, 1]), groups=self.in_channels, stride=2, padding=max(f_pad - w_pad, 0)) - x = torch.nn.functional.conv2d(x, w, padding=max(w_pad - f_pad, 0)) - elif self.fused_resample and self.down and w is not None: - x = torch.nn.functional.conv2d(x, w, padding=w_pad+f_pad) - x = torch.nn.functional.conv2d(x, f.tile([self.out_channels, 1, 1, 1]), groups=self.out_channels, stride=2) - else: - if self.up: - x = torch.nn.functional.conv_transpose2d(x, f.mul(4).tile([self.in_channels, 1, 1, 1]), groups=self.in_channels, stride=2, padding=f_pad) - if self.down: - x = torch.nn.functional.conv2d(x, f.tile([self.in_channels, 1, 1, 1]), groups=self.in_channels, stride=2, padding=f_pad) - if w is not None: - x = torch.nn.functional.conv2d(x, w, padding=w_pad) - if b is not None: - x = x.add_(b.reshape(1, -1, 1, 1)) - return x - - -class GroupNorm(torch.nn.Module): - def __init__(self, num_channels, num_groups=32, min_channels_per_group=4, eps=1e-5): - super().__init__() - self.num_groups = min(num_groups, num_channels // min_channels_per_group) - self.eps = eps - self.weight = torch.nn.Parameter(torch.ones(num_channels)) - self.bias = torch.nn.Parameter(torch.zeros(num_channels)) - - def forward(self, x): - x = torch.nn.functional.group_norm(x, num_groups=self.num_groups, weight=self.weight.to(x.dtype), bias=self.bias.to(x.dtype), eps=self.eps) - return x - -#---------------------------------------------------------------------------- -# Attention weight computation, i.e., softmax(Q^T * K). -# Performs all computation using FP32, but uses the original datatype for -# inputs/outputs/gradients to conserve memory. - -class AttentionOp(torch.autograd.Function): - @staticmethod - def forward(ctx, q, k): - w = torch.einsum('ncq,nck->nqk', q.to(torch.float32), (k / np.sqrt(k.shape[1])).to(torch.float32)).softmax(dim=2).to(q.dtype) - ctx.save_for_backward(q, k, w) - return w - - @staticmethod - def backward(ctx, dw): - q, k, w = ctx.saved_tensors - db = torch._softmax_backward_data(grad_output=dw.to(torch.float32), output=w.to(torch.float32), dim=2, input_dtype=torch.float32) - dq = torch.einsum('nck,nqk->ncq', k.to(torch.float32), db).to(q.dtype) / np.sqrt(k.shape[1]) - dk = torch.einsum('ncq,nqk->nck', q.to(torch.float32), db).to(k.dtype) / np.sqrt(k.shape[1]) - return dq, dk - - -class UNetBlock(torch.nn.Module): - def __init__(self, - in_channels, out_channels, emb_channels, up=False, down=False, attention=False, - num_heads=None, channels_per_head=64, dropout=0, skip_scale=1, eps=1e-5, - resample_filter=[1,1], resample_proj=False, adaptive_scale=True, - init=dict(), init_zero=dict(init_weight=0), init_attn=None, - ): - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.emb_channels = emb_channels - self.num_heads = 0 if not attention else num_heads if num_heads is not None else out_channels // channels_per_head - self.dropout = dropout - self.skip_scale = skip_scale - self.adaptive_scale = adaptive_scale - - self.norm0 = GroupNorm(num_channels=in_channels, eps=eps) - self.conv0 = Conv2d(in_channels=in_channels, out_channels=out_channels, kernel=3, up=up, down=down, resample_filter=resample_filter, **init) - self.affine = Linear(in_features=emb_channels, out_features=out_channels*(2 if adaptive_scale else 1), **init) - self.norm1 = GroupNorm(num_channels=out_channels, eps=eps) - self.conv1 = Conv2d(in_channels=out_channels, out_channels=out_channels, kernel=3, **init_zero) - - self.skip = None - if out_channels != in_channels or up or down: - kernel = 1 if resample_proj or out_channels!= in_channels else 0 - self.skip = Conv2d(in_channels=in_channels, out_channels=out_channels, kernel=kernel, up=up, down=down, resample_filter=resample_filter, **init) - - if self.num_heads: - self.norm2 = GroupNorm(num_channels=out_channels, eps=eps) - self.qkv = Conv2d(in_channels=out_channels, out_channels=out_channels*3, kernel=1, **(init_attn if init_attn is not None else init)) - self.proj = Conv2d(in_channels=out_channels, out_channels=out_channels, kernel=1, **init_zero) - - def forward(self, x, emb): - orig = x - x = self.conv0(silu(self.norm0(x))) - - params = self.affine(emb).unsqueeze(2).unsqueeze(3).to(x.dtype) - - # print(f'emb shape: {emb.shape}, affine out shape: {self.affine(emb).shape}, x shape: {x.shape}, params shape: {params.shape}') - - if self.adaptive_scale: - scale, shift = params.chunk(chunks=2, dim=1) - x = silu(torch.addcmul(shift, self.norm1(x), scale + 1)) - else: - x = silu(self.norm1(x.add_(params))) - - x = self.conv1(torch.nn.functional.dropout(x, p=self.dropout, training=self.training)) - x = x.add_(self.skip(orig) if self.skip is not None else orig) - x = x * self.skip_scale - - if self.num_heads: - q, k, v = self.qkv(self.norm2(x)).reshape(x.shape[0] * self.num_heads, x.shape[1] // self.num_heads, 3, -1).unbind(2) - w = AttentionOp.apply(q, k) - a = torch.einsum('nqk,nck->ncq', w, v) - x = self.proj(a.reshape(*x.shape)).add_(x) - x = x * self.skip_scale - return x - - -class UNet(torch.nn.Module): - def __init__(self, - img_resolution = 256, # Image resolution at input/output. - in_channels = 8, # Number of color channels at input. - out_channels = 8, # Number of color channels at output. - label_dim = 0, # Number of class labels, 0 = unconditional. - augment_dim = 0, # Augmentation label dimensionality, 0 = no augmentation. - - model_channels = 128, # Base multiplier for the number of channels. - channel_mult = [1,2,2,2], # Per-resolution multipliers for the number of channels. - channel_mult_emb = 4, # Multiplier for the dimensionality of the embedding vector. - num_blocks = 4, # Number of residual blocks per resolution. - attn_resolutions = [16], # List of resolutions with self-attention. - dropout = 0.10, # Dropout probability of intermediate activations. - label_dropout = 0, # Dropout probability of class labels for classifier-free guidance. - - embedding_type = 'positional', # Timestep embedding type: 'positional' for DDPM++, 'fourier' for NCSN++. - channel_mult_noise = 1, # Timestep embedding size: 1 for DDPM++, 2 for NCSN++. - encoder_type = 'standard', # Encoder architecture: 'standard' for DDPM++, 'residual' for NCSN++. - decoder_type = 'standard', # Decoder architecture: 'standard' for both DDPM++ and NCSN++. - resample_filter = [1,1], # Resampling filter: [1,1] for DDPM++, [1,3,3,1] for NCSN++. - ): - assert embedding_type in ['fourier', 'positional'] - assert encoder_type in ['standard', 'skip', 'residual'] - assert decoder_type in ['standard', 'skip'] - - super().__init__() - self.label_dropout = label_dropout - emb_channels = model_channels * channel_mult_emb - noise_channels = model_channels * channel_mult_noise - init = dict(init_mode='xavier_uniform') - init_zero = dict(init_mode='xavier_uniform', init_weight=1e-5) - init_attn = dict(init_mode='xavier_uniform', init_weight=np.sqrt(0.2)) - block_kwargs = dict( - emb_channels=emb_channels, num_heads=1, dropout=dropout, skip_scale=np.sqrt(0.5), eps=1e-6, - resample_filter=resample_filter, resample_proj=True, adaptive_scale=False, - init=init, init_zero=init_zero, init_attn=init_attn, - ) - - # Mapping. - self.map_noise = PositionalEmbedding(num_channels=noise_channels, endpoint=True) if embedding_type == 'positional' else FourierEmbedding(num_channels=noise_channels) - self.map_label = Linear(in_features=label_dim, out_features=noise_channels, **init) if label_dim else None - self.map_augment = Linear(in_features=augment_dim, out_features=noise_channels, bias=False, **init) if augment_dim else None - self.map_layer0 = Linear(in_features=noise_channels, out_features=emb_channels, **init) - self.map_layer1 = Linear(in_features=emb_channels, out_features=emb_channels, **init) - - # print(f"model_channels: {model_channels}, noise_channels: {noise_channels}, emb_channels: {emb_channels}") - # Encoder. - self.enc = torch.nn.ModuleDict() - cout = in_channels - caux = in_channels - for level, mult in enumerate(channel_mult): - res = img_resolution >> level - if level == 0: - cin = cout - cout = model_channels - self.enc[f'{res}x{res}_conv'] = Conv2d(in_channels=cin, out_channels=cout, kernel=3, **init) - else: - self.enc[f'{res}x{res}_down'] = UNetBlock(in_channels=cout, out_channels=cout, down=True, **block_kwargs) - if encoder_type == 'skip': - self.enc[f'{res}x{res}_aux_down'] = Conv2d(in_channels=caux, out_channels=caux, kernel=0, down=True, resample_filter=resample_filter) - self.enc[f'{res}x{res}_aux_skip'] = Conv2d(in_channels=caux, out_channels=cout, kernel=1, **init) - if encoder_type == 'residual': - self.enc[f'{res}x{res}_aux_residual'] = Conv2d(in_channels=caux, out_channels=cout, kernel=3, down=True, resample_filter=resample_filter, fused_resample=True, **init) - caux = cout - for idx in range(num_blocks): - cin = cout - cout = model_channels * mult - attn = (res in attn_resolutions) - self.enc[f'{res}x{res}_block{idx}'] = UNetBlock(in_channels=cin, out_channels=cout, attention=attn, **block_kwargs) - skips = [block.out_channels for name, block in self.enc.items() if 'aux' not in name] - - # Decoder. - self.dec = torch.nn.ModuleDict() - for level, mult in reversed(list(enumerate(channel_mult))): - res = img_resolution >> level - if level == len(channel_mult) - 1: - self.dec[f'{res}x{res}_in0'] = UNetBlock(in_channels=cout, out_channels=cout, attention=True, **block_kwargs) - self.dec[f'{res}x{res}_in1'] = UNetBlock(in_channels=cout, out_channels=cout, **block_kwargs) - else: - self.dec[f'{res}x{res}_up'] = UNetBlock(in_channels=cout, out_channels=cout, up=True, **block_kwargs) - for idx in range(num_blocks + 1): - cin = cout + skips.pop() - cout = model_channels * mult - attn = (idx == num_blocks and res in attn_resolutions) - self.dec[f'{res}x{res}_block{idx}'] = UNetBlock(in_channels=cin, out_channels=cout, attention=attn, **block_kwargs) - if decoder_type == 'skip' or level == 0: - if decoder_type == 'skip' and level < len(channel_mult) - 1: - self.dec[f'{res}x{res}_aux_up'] = Conv2d(in_channels=out_channels, out_channels=out_channels, kernel=0, up=True, resample_filter=resample_filter) - self.dec[f'{res}x{res}_aux_norm'] = GroupNorm(num_channels=cout, eps=1e-6) - self.dec[f'{res}x{res}_aux_conv'] = Conv2d(in_channels=cout, out_channels=out_channels, kernel=3, **init_zero) - - # def forward(self, x, noise_labels, class_labels, augment_labels=None): - def forward(self, x, field_labels, bcs, opts, augment_labels=None): - # Mapping. - # print(f'Input x shape: {x.shape}') - x = x.squeeze(0) - x = x.squeeze(2) # reshape from [1, B, C, 1, H, W] to [B, C, H, W] - # print(f'Reshaped x shape: {x.shape}') - noise_labels = getattr(opts, 'sigma', None) - diffusion_cond = getattr(opts, 'diffusion_cond', None) - if diffusion_cond is not None: - diffusion_cond = diffusion_cond.squeeze(1) - diffusion_cond = diffusion_cond.squeeze(2) # reshape from [B, 1, C, 1, H, W] to [B, C, H, W] - # print(f'Reshaped diffusion_cond shape: {diffusion_cond.shape}') - x = torch.cat([x, diffusion_cond], dim=1) # concatenate along channel dimension to get shape [B, 2*C, H, W] - # print(f'concatenated x shape: {x.shape}') - - if noise_labels is not None: - emb = self.map_noise(noise_labels) - emb = emb.reshape(emb.shape[0], 2, -1).flip(1).reshape(*emb.shape) # swap sin/cos - # if self.map_label is not None: - # class_labels = opts.diffusion_cond.reshape(opts.diffusion_cond.shape[0], -1) if hasattr(opts, 'diffusion_cond') else None - # tmp = class_labels - # if self.training and self.label_dropout: - # tmp = tmp * (torch.rand([x.shape[0], 1], device=x.device) >= self.label_dropout).to(tmp.dtype) - # emb = emb + self.map_label(tmp * np.sqrt(self.map_label.in_features)) - # if self.map_augment is not None and augment_labels is not None: - # emb = emb + self.map_augment(augment_labels) - - emb = silu(self.map_layer0(emb)) - emb = silu(self.map_layer1(emb)) - - # Encoder. - skips = [] - aux = x - for name, block in self.enc.items(): - if 'aux_down' in name: - aux = block(aux) - elif 'aux_skip' in name: - x = skips[-1] = x + block(aux) - elif 'aux_residual' in name: - x = skips[-1] = aux = (x + block(aux)) / np.sqrt(2) - else: - x = block(x, emb) if isinstance(block, UNetBlock) else block(x) - skips.append(x) - - # Decoder. - aux = None - tmp = None - for name, block in self.dec.items(): - if 'aux_up' in name: - aux = block(aux) - elif 'aux_norm' in name: - tmp = block(x) - elif 'aux_conv' in name: - tmp = block(silu(tmp)) - aux = tmp if aux is None else tmp + aux - else: - if x.shape[1] != block.in_channels: - x = torch.cat([x, skips.pop()], dim=1) - x = block(x, emb) - - if diffusion_cond is not None: - aux, _ = aux.chunk(2, dim=1) - # print(f'Output aux shape before reshape: {aux.shape}') - # print(f'Output aux shape after reshape: {aux.unsqueeze(2).shape}') - return aux.unsqueeze(2) # reshape from [B, C, H, W] to [B, C, 1, H, W] - - - - - - diff --git a/matey/models/vit.py b/matey/models/vit.py index ee82042..12b5095 100644 --- a/matey/models/vit.py +++ b/matey/models/vit.py @@ -7,8 +7,6 @@ import numpy as np from einops import rearrange from .spacetime_modules import SpaceTimeBlock_all2all -from .time_modules import FourierEmbedding, PositionalEmbedding, Linear -from torch.nn.functional import silu from .basemodel import BaseModel from ..data_utils.shared_utils import normalize_spatiotemporal_persample, get_top_variance_patchids, normalize_spatiotemporal_persample_graph from ..utils import ForwardOptionsBase, TrainOptionsBase, densenodes_to_graphnodes @@ -39,7 +37,6 @@ def build_vit(params): replace_patch=getattr(params, 'replace_patch', True), hierarchical=getattr(params, 'hierarchical', None), use_linear=getattr(params, 'use_linear', False), - diffusion=getattr(params, 'diffusion', False) ) return model @@ -56,7 +53,7 @@ class ViT_all2all(BaseModel): sts_f """ def __init__(self, tokenizer_heads=None, embed_dim=768, 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", replace_patch=True, SR_ratio=[1,1,1], hierarchical=None, use_linear=False, diffusion=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, use_linear=False): 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, use_linear=use_linear) self.drop_path = drop_path @@ -74,33 +71,6 @@ def __init__(self, tokenizer_heads=None, embed_dim=768, num_heads=12, processor self.processor_blocks=processor_blocks self.replace_patch=replace_patch assert not (self.replace_patch and self.sts_model) - - self.diffusion=diffusion - if self.diffusion: - #from https://github.com/NVlabs/edm/blob/008a4e5316c8e3bfe61a62f874bddba254295afb/training/networks.py#L269 - #FIXME: @Paul, currently place holder pls check the proper setting of these variables - model_channels = 128 # Base multiplier for the number of channels. - channel_mult_emb = 4 # Multiplier for the dimensionality of the embedding vector. - embedding_type = 'positional' # Timestep embedding type: 'positional' for DDPM++, 'fourier' for NCSN++. - channel_mult_noise = 1 # Timestep embedding size: 1 for DDPM++, 2 for NCSN++. - emb_channels = model_channels * channel_mult_emb - noise_channels = model_channels * channel_mult_noise - init = dict(init_mode='xavier_uniform') - - self.map_noise = PositionalEmbedding(num_channels=noise_channels, endpoint=True) if embedding_type == 'positional' else FourierEmbedding(num_channels=noise_channels) - self.map_layer0 = Linear(in_features=noise_channels, out_features=emb_channels, **init) - # self.map_layer1 = Linear(in_features=emb_channels, out_features=embed_dim, **init) - self.map_layer1 = Linear(in_features=emb_channels, out_features=emb_channels, **init) - - # self.affine = nn.ModuleDict({}) - # for imod in range(self.nhlevels): - # self.affine[str(imod)] = Linear(in_features=emb_channels, out_features=embed_dim, **init) - self.affine = Linear(in_features=emb_channels, out_features=embed_dim, **init) - - # self.skip_projection = nn.ModuleDict({}) - # for imod in range(self.nhlevels): - # self.skip_projection[str(imod)] = Linear(in_features=embed_dim*2, out_features=embed_dim, **init) - self.skip_projection = Linear(in_features=embed_dim*2, out_features=embed_dim, **init) def expand_sts_model(self): """ Appends addition sts blocks""" @@ -161,27 +131,11 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase, train_opts: cond_input = opts.cond_input isgraph=opts.isgraph field_labels_out=opts.field_labels_out - sigma = getattr(opts, 'sigma', None) - diffusion_cond = getattr(opts, 'diffusion_cond', None) ################################################################## 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.map_noise(sigma) - # print(f'noise_labels shape: {sigma.shape}, emb shape: {emb.shape}') - emb = emb.reshape(emb.shape[0], 2, -1).flip(1).reshape(*emb.shape) # swap sin/cos - # print(f'after swap emb shape: {emb.shape}') - emb = silu(self.map_layer0(emb)) - # print(f'after first layer emb shape: {emb.shape}') - emb = silu(self.map_layer1(emb)) #in shape - # print(f'after second layer emb shape: {emb.shape}') - emb = self.affine(emb) - # print(f'after affine emb shape: {emb.shape}') - 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: x = data.x#[nnodes, T, C] @@ -220,13 +174,6 @@ 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, _, _, _, _, _, _, _ = self.get_patchsequence(diffusion_cond, state_labels, tkhead_name, refineind=refineind, blockdict=blockdict, ilevel=imod, isgraph=isgraph) - diffusion_cond = rearrange(diffusion_cond, 't b c ntoken_tot -> b c (t ntoken_tot)') - - # if self.diffusion: - # x_padding = x_padding + emb.unsqueeze(-1) - # Repeat the steps for conditioning if present if conditioning: assert self.sts_model == False @@ -247,35 +194,11 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase, train_opts: x_padding = x_padding + c if iblk==0: - if self.diffusion: - if diffusion_cond is not None: - # add diffusion_cond as additional tokens for the diffusion model to attend to - x_padding = torch.cat([x_padding, diffusion_cond], dim=2) - - # add noise embedding to input tokens as conditional information for diffusion model - x_padding = torch.cat([x_padding, emb.unsqueeze(-1)], dim=2) - - x_input = x_padding.clone() - x_padding = blk(x_padding, sequence_parallel_group=sequence_parallel_group, bcs=bcs, leadtime=leadtime, mask_padding=mask4attblk ) else: - if iblk==len(self.blocks)-1 and self.diffusion: - # for the last block, add skip connection from input tokens to facilitate training of diffusion model - x_padding = torch.cat([x_padding, x_input], dim=1) - local_batch = x_padding.shape[0] - x_padding = rearrange(x_padding, 'b c L -> (b L) c') - x_padding = self.skip_projection(x_padding) # Residual connection from input to the last block in the level - x_padding = rearrange(x_padding, '(b L) c -> b c L', b = local_batch) x_padding = blk(x_padding, sequence_parallel_group=sequence_parallel_group, bcs=bcs, leadtime=None, mask_padding=mask4attblk) #self.debug_nan(x_padding, message="attention block") ################################################################################ - - if self.diffusion: - x_padding = x_padding[:, :, :-1] # remove noise embedding from tokens after processing - if diffusion_cond is not None: - doubled_L = x_padding.shape[2] - x_padding = x_padding[:, :, :doubled_L//2] # remove diffusion_cond part from tokens after processing - x_padding = rearrange(x_padding, 'b c (t ntoken_tot) -> t b c ntoken_tot', t=T) ######## Decode ######## if self.sts_model: diff --git a/matey/train.py b/matey/train.py index 29eb7ef..3d70d34 100644 --- a/matey/train.py +++ b/matey/train.py @@ -21,7 +21,6 @@ from .models.svit import build_svit from .models.vit import build_vit from .models.turbt import build_turbt -from .models.turbt_modified import build_turbt_modified 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 @@ -194,8 +193,6 @@ def initialize_model(self): self.model = build_vit(self.params).to(self.device) elif self.params.model_type == "turbt": self.model = build_turbt(self.params).to(self.device) - elif self.params.model_type == "turbt_modified": - self.model = build_turbt_modified(self.params).to(self.device) if self.params.compile: print('WARNING: BFLOAT NOT SUPPORTED IN SOME COMPILE OPS SO SWITCHING TO FLOAT16') From f383a54d6e3ed66447691ddf272fca49acacdaa3 Mon Sep 17 00:00:00 2001 From: Paul Laiu Date: Wed, 20 May 2026 12:53:49 -0400 Subject: [PATCH 06/12] moved diffusion model specific architectures to the base model; fixed an unused import in diffusion_model.py --- examples/config/Demo_MW_diffusion_TT.yaml | 6 +++- examples/submit_batch_diffusion.sh | 4 +-- matey/models/basemodel.py | 27 ++++++++++++++-- matey/models/diffusion_model.py | 1 - matey/models/turbt.py | 38 ++++++----------------- 5 files changed, 41 insertions(+), 35 deletions(-) diff --git a/examples/config/Demo_MW_diffusion_TT.yaml b/examples/config/Demo_MW_diffusion_TT.yaml index 015adad..00ca247 100644 --- a/examples/config/Demo_MW_diffusion_TT.yaml +++ b/examples/config/Demo_MW_diffusion_TT.yaml @@ -8,7 +8,7 @@ basic_config: &basic_config 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_TurbT_3level_S_lt5' # Output path + 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 @@ -46,6 +46,10 @@ basic_config: &basic_config processor_blocks: 12 # Number of transformer blocks in the backbone - 12/12/12/24 for Ti/S/B/L 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]] diff --git a/examples/submit_batch_diffusion.sh b/examples/submit_batch_diffusion.sh index 8d2b037..0e2f775 100644 --- a/examples/submit_batch_diffusion.sh +++ b/examples/submit_batch_diffusion.sh @@ -2,10 +2,10 @@ #SBATCH -A LRN037 #SBATCH -J matey #SBATCH -o %x-%j.out -#SBATCH -t 02:00:00 +#SBATCH -t 00:20:00 #SBATCH -p batch #SBATCH -N 1 -##SBATCH -q debug +#SBATCH -q debug #SBATCH -C nvme export OMP_NUM_THREADS=1 diff --git a/matey/models/basemodel.py b/matey/models/basemodel.py index 95fa091..803f6fb 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, PositionalEmbedding, Linear from .input_modules import input_control_MLP from .positionbias_modules import positionbias_mod import sys @@ -28,7 +28,8 @@ class BaseModel(nn.Module): embed_dim (int): Dimension of the embedding 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): + 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, + diffusion=False, model_channels=128, channel_mult_emb=4, embedding_type='positional', channel_mult_noise=1): 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) @@ -90,7 +91,27 @@ 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 + self.diffusion = diffusion + if self.diffusion: + emb_channels = model_channels * channel_mult_emb + noise_channels = model_channels * channel_mult_noise + init = dict(init_mode='xavier_uniform') + + self.map_noise = (PositionalEmbedding(num_channels=noise_channels, endpoint=True) + if embedding_type == 'positional' + else FourierEmbedding(num_channels=noise_channels)) + self.map_layer0 = Linear(in_features=noise_channels, out_features=emb_channels, **init) + self.map_layer1 = Linear(in_features=emb_channels, out_features=emb_channels, **init) + + self.affine = nn.ModuleDict({}) + for ilevel in range(nlevels): + self.affine[str(ilevel)] = Linear(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(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""" diff --git a/matey/models/diffusion_model.py b/matey/models/diffusion_model.py index 3f8ef66..5c67b71 100644 --- a/matey/models/diffusion_model.py +++ b/matey/models/diffusion_model.py @@ -1,4 +1,3 @@ -from matey.models.unet import build_unet import torch import torch.nn as nn diff --git a/matey/models/turbt.py b/matey/models/turbt.py index a394e88..3898b76 100644 --- a/matey/models/turbt.py +++ b/matey/models/turbt.py @@ -12,7 +12,6 @@ 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 -from .time_modules import FourierEmbedding, PositionalEmbedding, Linear from torch.nn.functional import silu import sys, copy from operator import mul @@ -45,6 +44,10 @@ def build_turbt(params): hierarchical=getattr(params, 'hierarchical', None), notransposed=getattr(params, 'notransposed', False), diffusion=getattr(params, 'diffusion', False), + model_channels=getattr(params, 'model_channels', 128), + channel_mult_emb=getattr(params, 'channel_mult_emb', 4), + embedding_type=getattr(params, 'embedding_type', 'positional'), + channel_mult_noise=getattr(params, 'channel_mult_noise', 1), ) return model @@ -61,9 +64,12 @@ class TurbT(BaseModel): sts_f """ def __init__(self, tokenizer_heads=None, embed_dim=768, num_heads=12, processor_blocks=8, n_states=6, - 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=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) + 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=False, model_channels=128, channel_mult_emb=4, embedding_type='positional', channel_mult_noise=1): + 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, + diffusion=diffusion, model_channels=model_channels, channel_mult_emb=channel_mult_emb, + embedding_type=embedding_type, channel_mult_noise=channel_mult_noise) self.drop_path = drop_path self.dp = np.linspace(0, drop_path, processor_blocks) self.module_blocks=nn.ModuleDict({}) @@ -120,30 +126,6 @@ def __init__(self, tokenizer_heads=None, embed_dim=768, num_heads=12, processor 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)]) - self.diffusion = diffusion - if self.diffusion: - model_channels = 128 - channel_mult_emb = 4 - embedding_type = 'positional' - 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(num_channels=noise_channels, endpoint=True) - if embedding_type == 'positional' - else FourierEmbedding(num_channels=noise_channels)) - self.map_layer0 = Linear(in_features=noise_channels, out_features=emb_channels, **init) - self.map_layer1 = Linear(in_features=emb_channels, out_features=emb_channels, **init) - - self.affine = nn.ModuleDict({}) - for imod in range(self.nhlevels): - self.affine[str(imod)] = Linear(in_features=emb_channels, out_features=embed_dim, **init) - - self.diffusion_cond_proj = nn.ModuleDict({}) - for imod in range(self.nhlevels): - self.diffusion_cond_proj[str(imod)] = Linear(in_features=embed_dim, out_features=embed_dim, **init) - 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}" From 098857cd23c4c3bd45456593dfc9fae39d31435e Mon Sep 17 00:00:00 2001 From: Paul Laiu Date: Wed, 20 May 2026 15:57:16 -0400 Subject: [PATCH 07/12] added diffusion model options to vit, avit, and svit; added associated config files --- examples/config/Demo_MW_diffusion_avit.yaml | 87 +++++++ examples/config/Demo_MW_diffusion_svit.yaml | 87 +++++++ ...on_ViT.yaml => Demo_MW_diffusion_vit.yaml} | 6 +- examples/submit_batch_diffusion.sh | 8 +- examples/submit_batch_generate.sh | 13 +- matey/generate.py | 234 ++++++++---------- matey/models/attention_modules.py | 4 + matey/models/avit.py | 53 +++- matey/models/svit.py | 52 +++- matey/models/vit.py | 50 +++- matey/train.py | 4 +- 11 files changed, 435 insertions(+), 163 deletions(-) create mode 100644 examples/config/Demo_MW_diffusion_avit.yaml create mode 100644 examples/config/Demo_MW_diffusion_svit.yaml rename examples/config/{Demo_MW_diffusion_ViT.yaml => Demo_MW_diffusion_vit.yaml} (92%) diff --git a/examples/config/Demo_MW_diffusion_avit.yaml b/examples/config/Demo_MW_diffusion_avit.yaml new file mode 100644 index 0000000..1b9d5d6 --- /dev/null +++ b/examples/config/Demo_MW_diffusion_avit.yaml @@ -0,0 +1,87 @@ +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: !!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..4896658 --- /dev/null +++ b/examples/config/Demo_MW_diffusion_svit.yaml @@ -0,0 +1,87 @@ +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: !!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 similarity index 92% rename from examples/config/Demo_MW_diffusion_ViT.yaml rename to examples/config/Demo_MW_diffusion_vit.yaml index 75d3539..f99c8e7 100644 --- a/examples/config/Demo_MW_diffusion_ViT.yaml +++ b/examples/config/Demo_MW_diffusion_vit.yaml @@ -8,7 +8,7 @@ basic_config: &basic_config 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_newembskip_lt5' # Output path + 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 @@ -46,6 +46,10 @@ basic_config: &basic_config processor_blocks: 12 # Number of transformer blocks in the backbone - 12/12/12/24 for Ti/S/B/L 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]] diff --git a/examples/submit_batch_diffusion.sh b/examples/submit_batch_diffusion.sh index 0e2f775..ad75d0e 100644 --- a/examples/submit_batch_diffusion.sh +++ b/examples/submit_batch_diffusion.sh @@ -2,10 +2,10 @@ #SBATCH -A LRN037 #SBATCH -J matey #SBATCH -o %x-%j.out -#SBATCH -t 00:20:00 +#SBATCH -t 01:00:00 #SBATCH -p batch #SBATCH -N 1 -#SBATCH -q debug +##SBATCH -q debug #SBATCH -C nvme export OMP_NUM_THREADS=1 @@ -14,8 +14,10 @@ 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 -# export yaml_config=./config/Demo_MW_diffusion_ViT.yaml ##conda env with rocm 6.0.0 diff --git a/examples/submit_batch_generate.sh b/examples/submit_batch_generate.sh index 6a329a9..dcb2e2a 100644 --- a/examples/submit_batch_generate.sh +++ b/examples/submit_batch_generate.sh @@ -13,10 +13,15 @@ export OMP_NUM_THREADS=1 export master_node=$SLURMD_NODENAME export config="basic_config" - -export model_dir="./Demo_Diffusion_MW_cond_TurbT_3level_S_lt5/basic_config/demo_diffusion/" - -export output_dir="./MW_cond_generation_outputs_batches_lt5/turbt/" +# 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 diff --git a/matey/generate.py b/matey/generate.py index ccd1411..ff47241 100644 --- a/matey/generate.py +++ b/matey/generate.py @@ -96,7 +96,7 @@ def initialize_data(self): self.single_print("self.train_data_loader:", len(self.train_data_loader), "valid_data_loader:", len(self.valid_data_loader)) def initialize_model(self): - if self.params.diffusion: + if getattr(self.params, "diffusion", False): self.model = build_diffusion_model(self.params).to(self.device) else: raise NotImplementedError("Only diffusion model generation is implemented currently. Please set params.diffusion to True.") @@ -215,29 +215,29 @@ def generate(self, seed=None, num_samples=1, batch_list=[0]): sigma_min = 0.002 sigma_max = 80 - num_steps = 18 + n_diffusion_steps = 18 rho=7 S_churn=0 S_min=0 S_max=float('inf') S_noise=1 - step_indices = torch.arange(num_steps, device=self.device) + step_indices = torch.arange(n_diffusion_steps, device=self.device) - t_steps = (sigma_max ** (1 / rho) + step_indices / (num_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho + t_steps = (sigma_max ** (1 / rho) + step_indices / (n_diffusion_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho t_steps = torch.cat([t_steps, torch.zeros_like(t_steps[:1])]) # t_N = 0 # tar = tar.to(self.device) # init_inp = rearrange(init_inp.to(self.device), 'b t c d h w -> t b c d h w') - + x_next = init_inp * t_steps[0] for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])): # 0, ..., N-1 x_cur = x_next # Increase noise temporarily. - gamma = min(S_churn / num_steps, np.sqrt(2) - 1) if S_min <= t_cur <= S_max else 0 + gamma = min(S_churn / n_diffusion_steps, np.sqrt(2) - 1) if S_min <= t_cur <= S_max else 0 t_hat = t_cur + gamma * t_cur x_hat = x_cur + (t_hat ** 2 - t_cur ** 2).sqrt() * S_noise * torch.randn_like(x_cur) @@ -254,7 +254,7 @@ def generate(self, seed=None, num_samples=1, batch_list=[0]): isgraph=isgraph, field_labels_out= field_labels ) - if self.cond_diffusion: + if getattr(self.params, "cond_diffusion", False): opts.diffusion_cond = rearrange(inp.to(self.device), 't b c d h w -> b t c d h w') # conditioning on input history # Euler step. # denoised = net(x_hat, t_hat, class_labels).to(torch.float64) @@ -269,7 +269,7 @@ def generate(self, seed=None, num_samples=1, batch_list=[0]): # print(f"x_next shape: {x_next.shape}, d_cur shape: {d_cur.shape}, t_next shape: {t_next.shape}, t_hat shape: {t_hat.shape}") # Apply 2nd order correction. - if i < num_steps - 1: + if i < n_diffusion_steps - 1: opts = ForwardOptionsBase( imod=imod, @@ -283,7 +283,7 @@ def generate(self, seed=None, num_samples=1, batch_list=[0]): isgraph=isgraph, field_labels_out= field_labels ) - if self.cond_diffusion: + if getattr(self.params, "cond_diffusion", False): opts.diffusion_cond = rearrange(inp.to(self.device), 't b c d h w -> b t c d h w') # conditioning on input history # denoised = net(x_next, t_next, class_labels).to(torch.float64) # denoised = net(x_next, t_next, None) @@ -310,168 +310,140 @@ def generate(self, seed=None, num_samples=1, batch_list=[0]): 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 Generation Loop...") + self.single_print("Starting Autoregressive Generation Loop...") if seed is not None: - seed_value = seed - torch.manual_seed(seed_value) + torch.manual_seed(seed) + # Load one batch of data data_iter = iter(self.valid_data_loader) - for step_idx in range(num_steps): - - - if step_idx==0: - data = next(data_iter) - - if "graph" in data: - graphdata = data["graph"].to(self.device) - tar = graphdata.y #[nnodes, C_tar] - leadtime = graphdata.leadtime #[nnodes, 1] - 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 - if supportdata: - cond_input = data["cond_input"].to(self.device) - else: - cond_input = 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 - - 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")}') + data = next(data_iter) - for sample_idx in range(num_samples): + 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 - - init_inp = torch.randn(inp.shape).to(self.device) + seq_group = self.current_group if dset_type in self.valid_dataset.DP_dsets else None - - self.model.eval() + # 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")}') - with torch.no_grad(): + self.model.eval() + inp_cur = inp.clone() # conditioning that advances each autoregressive step + for step_idx in range(num_steps): + samples = [] # collect num_samples independent draws at this step + for sample_idx in range(num_samples): + with torch.no_grad(): sigma_min = 0.002 sigma_max = 80 - num_steps = 18 - rho=7 - S_churn=0 - S_min=0 - S_max=float('inf') - S_noise=1 - - step_indices = torch.arange(num_steps, device=self.device) - - t_steps = (sigma_max ** (1 / rho) + step_indices / (num_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho - t_steps = torch.cat([t_steps, torch.zeros_like(t_steps[:1])]) # t_N = 0 + n_diffusion_steps = 18 + rho = 7 + S_churn = 0 + S_min = 0 + S_max = float('inf') + S_noise = 1 - # tar = tar.to(self.device) - # init_inp = rearrange(init_inp.to(self.device), 'b t c d h w -> t b c d h w') + step_indices = torch.arange(n_diffusion_steps, device=self.device) + t_steps = (sigma_max ** (1 / rho) + step_indices / (n_diffusion_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho + t_steps = torch.cat([t_steps, torch.zeros_like(t_steps[:1])]) # t_N = 0 - - x_next = init_inp * t_steps[0] + x_next = torch.randn(inp_cur.shape, device=self.device) * t_steps[0] - for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])): # 0, ..., N-1 + for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])): x_cur = x_next # Increase noise temporarily. - gamma = min(S_churn / num_steps, np.sqrt(2) - 1) if S_min <= t_cur <= S_max else 0 + gamma = min(S_churn / n_diffusion_steps, np.sqrt(2) - 1) if S_min <= t_cur <= S_max else 0 t_hat = t_cur + gamma * t_cur x_hat = x_cur + (t_hat ** 2 - t_cur ** 2).sqrt() * S_noise * torch.randn_like(x_cur) - opts = ForwardOptionsBase( - imod=imod, - imod_bottom=imod_bottom , - tkhead_name=tkhead_name, - sequence_parallel_group=seq_group, - leadtime=leadtime, - blockdict=copy.deepcopy(blockdict), - cond_dict=copy.deepcopy(cond_dict), - cond_input=cond_input, - isgraph=isgraph, - field_labels_out= field_labels + imod=imod, + imod_bottom=imod_bottom, + tkhead_name=tkhead_name, + sequence_parallel_group=seq_group, + leadtime=leadtime, + blockdict=copy.deepcopy(blockdict), + cond_dict=copy.deepcopy(cond_dict), + cond_input=cond_input, + isgraph=isgraph, + field_labels_out=field_labels ) if self.cond_diffusion: - opts.diffusion_cond = rearrange(inp.to(self.device), 't b c d h w -> b t c d h w') # conditioning on input history + opts.diffusion_cond = rearrange(inp_cur, 't b c d h w -> b t c d h w') + # Euler step. - # denoised = net(x_hat, t_hat, class_labels).to(torch.float64) - # denoised = net(x_hat, t_hat, None) - - # denoised = self.inference(x_hat, t_hat, field_labels, bcs, opts) denoised = self.model(x_hat, t_hat.repeat(x_hat.shape[1]), field_labels, bcs, opts) - - # print(f"denoised shape: {denoised.shape}, x_hat shape: {x_hat.shape}, t_hat shape: {t_hat.shape}") d_cur = (x_hat - denoised) / t_hat x_next = x_hat + (t_next - t_hat) * d_cur - # print(f"x_next shape: {x_next.shape}, d_cur shape: {d_cur.shape}, t_next shape: {t_next.shape}, t_hat shape: {t_hat.shape}") # Apply 2nd order correction. - if i < num_steps - 1: - + if i < n_diffusion_steps - 1: opts = ForwardOptionsBase( - imod=imod, - imod_bottom=imod_bottom , - tkhead_name=tkhead_name, - sequence_parallel_group=seq_group, - leadtime=leadtime, - blockdict=copy.deepcopy(blockdict), - cond_dict=copy.deepcopy(cond_dict), - cond_input=cond_input, - isgraph=isgraph, - field_labels_out= field_labels + imod=imod, + imod_bottom=imod_bottom, + tkhead_name=tkhead_name, + sequence_parallel_group=seq_group, + leadtime=leadtime, + blockdict=copy.deepcopy(blockdict), + cond_dict=copy.deepcopy(cond_dict), + cond_input=cond_input, + isgraph=isgraph, + field_labels_out=field_labels ) if self.cond_diffusion: - opts.diffusion_cond = rearrange(inp.to(self.device), 't b c d h w -> b t c d h w') # conditioning on input history - # denoised = net(x_next, t_next, class_labels).to(torch.float64) - # denoised = net(x_next, t_next, None) - # denoised = self.inference(x_next, t_next, field_labels, bcs, opts) - denoised = self.model(x_next, t_next.repeat(x_next.shape[1]), field_labels, bcs, opts) + opts.diffusion_cond = rearrange(inp_cur, 't b c d h w -> b t c d h w') + denoised = self.model(x_next, t_next.repeat(x_next.shape[1]), field_labels, bcs, opts) d_prime = (x_next - denoised) / t_next x_next = x_hat + (t_next - t_hat) * (0.5 * d_cur + 0.5 * d_prime) - # torch.save(x_next.cpu(), os.path.join(self.output_dir, f'generation_output_step_{i}_batch_{batch_idx}_sample{sample_idx}.pt')) - - output = x_next - - ###full resolution### - # residuals = output - tar - # torch.save(output.cpu(), os.path.join('Demo_Diffusion_CIFAR10_finemodel/basic_config/demo_diffusion/training_checkpoints/', 'generation_output.pt')) - # self.single_print(f'Generation output saved to {"Demo_Diffusion_CIFAR10_finemodel/basic_config/demo_diffusion/training_checkpoints/generation_output.pt"}') - torch.save(output.cpu(), os.path.join(self.output_dir, f'generation_output_batch_{batch_idx}_sample{sample_idx}.pt')) - self.single_print(f'Generation output saved to {os.path.join(self.output_dir, f"generation_output_batch_{batch_idx}_sample{sample_idx}.pt")}') + samples.append(x_next) + torch.save(x_next.cpu(), os.path.join(self.output_dir, f'generation_output_step_{step_idx}_sample{sample_idx}.pt')) + self.single_print(f'Generation output saved to {os.path.join(self.output_dir, f"generation_output_step_{step_idx}_sample{sample_idx}.pt")}') + # Mean of all samples becomes the conditioning for the next autoregressive step + inp_cur = torch.stack(samples, dim=0).mean(dim=0) + 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() + torch.cuda.empty_cache() diff --git a/matey/models/attention_modules.py b/matey/models/attention_modules.py index 584be0e..8bf771a 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..b6f15a5 100644 --- a/matey/models/avit.py +++ b/matey/models/avit.py @@ -7,6 +7,7 @@ import torch.nn.functional as F import numpy as np from einops import rearrange, repeat +from torch.nn.functional import silu from .spacetime_modules import SpaceTimeBlock from .basemodel import BaseModel from ..data_utils.shared_utils import normalize_spatiotemporal_persample, get_top_variance_patchids @@ -36,7 +37,12 @@ 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=getattr(params, 'diffusion', False), + model_channels=getattr(params, 'model_channels', 128), + channel_mult_emb=getattr(params, 'channel_mult_emb', 4), + embedding_type=getattr(params, 'embedding_type', 'positional'), + channel_mult_noise=getattr(params, 'channel_mult_noise', 1), ) return model @@ -52,9 +58,12 @@ 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=False, model_channels=128, channel_mult_emb=4, embedding_type='positional', channel_mult_noise=1): 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=diffusion, model_channels=model_channels, channel_mult_emb=channel_mult_emb, + embedding_type=embedding_type, channel_mult_noise=channel_mult_noise) self.drop_path = drop_path self.dp = np.linspace(0, drop_path, processor_blocks) @@ -165,16 +174,29 @@ 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.map_noise(sigma) + emb = emb.reshape(emb.shape[0], 2, -1).flip(1).reshape(*emb.shape) + emb = silu(self.map_layer0(emb)) + emb = silu(self.map_layer1(emb)) + emb = self.affine[str(imod)](emb) # (B, embed_dim) + 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 +224,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 +265,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/svit.py b/matey/models/svit.py index c5ccbdc..d0569d1 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 torch.nn.functional import silu 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 @@ -37,7 +38,12 @@ def build_svit(params): n_steps=params.n_steps, bias_type=params.bias_type, replace_patch=getattr(params, 'replace_patch', True), - hierarchical=getattr(params, 'hierarchical', None) + hierarchical=getattr(params, 'hierarchical', None), + diffusion=getattr(params, 'diffusion', False), + model_channels=getattr(params, 'model_channels', 128), + channel_mult_emb=getattr(params, 'channel_mult_emb', 4), + embedding_type=getattr(params, 'embedding_type', 'positional'), + channel_mult_noise=getattr(params, 'channel_mult_noise', 1), ) return model @@ -52,9 +58,12 @@ 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, - 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) + 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=False, model_channels=128, channel_mult_emb=4, embedding_type='positional', channel_mult_noise=1): + 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, + diffusion=diffusion, model_channels=model_channels, channel_mult_emb=channel_mult_emb, + embedding_type=embedding_type, channel_mult_noise=channel_mult_noise) self.drop_path = drop_path self.dp = np.linspace(0, drop_path, processor_blocks) @@ -133,12 +142,24 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase, train_opts: cond_input = opts.cond_input isgraph=opts.isgraph field_labels_out=opts.field_labels_out + 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.map_noise(sigma) + emb = emb.reshape(emb.shape[0], 2, -1).flip(1).reshape(*emb.shape) + emb = silu(self.map_layer0(emb)) + emb = silu(self.map_layer1(emb)) + emb = self.affine[str(imod)](emb) # (B, embed_dim) + 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: x = data.x#nnodes, T, C edge_index = data.edge_index # @@ -154,8 +175,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) @@ -179,6 +200,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 @@ -223,10 +257,12 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase, train_opts: return x[:, -1, :] #[nnodes, C] ######### 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/vit.py b/matey/models/vit.py index 12b5095..9bee6e8 100644 --- a/matey/models/vit.py +++ b/matey/models/vit.py @@ -6,6 +6,7 @@ import torch.nn as nn import numpy as np from einops import rearrange +from torch.nn.functional import silu from .spacetime_modules import SpaceTimeBlock_all2all from .basemodel import BaseModel from ..data_utils.shared_utils import normalize_spatiotemporal_persample, get_top_variance_patchids, normalize_spatiotemporal_persample_graph @@ -37,6 +38,11 @@ def build_vit(params): replace_patch=getattr(params, 'replace_patch', True), hierarchical=getattr(params, 'hierarchical', None), use_linear=getattr(params, 'use_linear', False), + diffusion=getattr(params, 'diffusion', False), + model_channels=getattr(params, 'model_channels', 128), + channel_mult_emb=getattr(params, 'channel_mult_emb', 4), + embedding_type=getattr(params, 'embedding_type', 'positional'), + channel_mult_noise=getattr(params, 'channel_mult_noise', 1), ) return model @@ -53,9 +59,12 @@ class ViT_all2all(BaseModel): sts_f """ def __init__(self, tokenizer_heads=None, embed_dim=768, 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", replace_patch=True, SR_ratio=[1,1,1], hierarchical=None, use_linear=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, use_linear=False, + diffusion=False, model_channels=128, channel_mult_emb=4, embedding_type='positional', channel_mult_noise=1): 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, use_linear=use_linear) + cond_input=cond_input, n_steps=n_steps, bias_type=bias_type, SR_ratio=SR_ratio, hierarchical=hierarchical, use_linear=use_linear, + diffusion=diffusion, model_channels=model_channels, channel_mult_emb=channel_mult_emb, + embedding_type=embedding_type, channel_mult_noise=channel_mult_noise) self.drop_path = drop_path self.dp = np.linspace(0, drop_path, processor_blocks) self.blocks = nn.ModuleList([SpaceTimeBlock_all2all(embed_dim, num_heads,drop_path=self.dp[i]) @@ -131,12 +140,24 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase, train_opts: cond_input = opts.cond_input isgraph=opts.isgraph field_labels_out=opts.field_labels_out + 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.map_noise(sigma) + emb = emb.reshape(emb.shape[0], 2, -1).flip(1).reshape(*emb.shape) + emb = silu(self.map_layer0(emb)) + emb = silu(self.map_layer1(emb)) + emb = self.affine[str(imod)](emb) # (B, embed_dim) + 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: x = data.x#[nnodes, T, C] edge_index = data.edge_index # @@ -153,9 +174,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) @@ -174,6 +194,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 @@ -230,9 +264,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 3d70d34..f14bd29 100644 --- a/matey/train.py +++ b/matey/train.py @@ -94,7 +94,7 @@ 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.params.diffusion: + if getattr(self.params, "diffusion", False): self.diffusion_loss = EDMLoss() self.initialize_model() @@ -183,7 +183,7 @@ def initialize_data(self): self.val_sampler.set_epoch(0) def initialize_model(self): - if self.params.diffusion: + if getattr(self.params, "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) From b716e0fca87756b71ba34c4031fe260c26551fd0 Mon Sep 17 00:00:00 2001 From: Paul Laiu Date: Wed, 20 May 2026 16:20:31 -0400 Subject: [PATCH 08/12] updated training_utils.py to match current version --- matey/utils/training_utils.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/matey/utils/training_utils.py b/matey/utils/training_utils.py index d35e6ef..f077dee 100644 --- a/matey/utils/training_utils.py +++ b/matey/utils/training_utils.py @@ -11,7 +11,6 @@ from .visualization_utils import checking_data_pred_tar import copy - def preprocess_target(leadtime, ramping_warmup = False): """ #Inputs: @@ -60,7 +59,13 @@ def autoregressive_rollout(model, inp, field_labels, bcs, opts: ForwardOptionsBa # output: Model output after the final autoregressive step ([B, C, D, H, W]) # rollout_steps: Number of autoregressive steps performed. """ - rollout_steps = preprocess_target(opts.leadtime) + is_constant = torch.all(opts.leadtime == opts.leadtime[0, 0]) + if is_constant: + rollout_steps = int(opts.leadtime[0,0].item()) + else: + rollout_steps = preprocess_target(opts.leadtime) + raise ValueError(f"Not expecting unequal leadtime across samples, {rollout_steps, opts.leadtime, is_constant}") + x_t = inp if opts.isgraph: n_steps = x_t.x.shape[1] #[nnodes, T, C] @@ -153,6 +158,7 @@ def compute_loss_and_logs(output, tar, graphdata, logs, loss_logs, dset_type, pa logs['train_nrmse'] += log_nrmse loss_logs[dset_type] += loss.item() logs['train_rmse'] += residuals.pow(2).mean(spatial_dims).sqrt().mean() + return loss, log_nrmse def update_loss_logs_inplace_eval(output, tar, graphdata, logs, loss_dset_logs, loss_l1_dset_logs, loss_rmse_dset_logs, dset_type): @@ -208,4 +214,3 @@ def __call__(self, net, images, field_labels, bcs, opts, augment_pipe=None): loss = weight * ((D_yn - y) ** 2) return loss, D_yn -#---------------------------------------------------------------------------- From ad18ddaf5d85581555ca7be108fa9fc524d12cfb Mon Sep 17 00:00:00 2001 From: Paul Laiu Date: Thu, 21 May 2026 16:16:52 -0400 Subject: [PATCH 09/12] updated handling of diffusion configurations; uncommented the model directory setting in submit_batch_generate.sh --- examples/config/Demo_MW_diffusion_TT.yaml | 13 ++--- examples/config/Demo_MW_diffusion_avit.yaml | 13 ++--- examples/config/Demo_MW_diffusion_svit.yaml | 13 ++--- examples/config/Demo_MW_diffusion_vit.yaml | 13 ++--- examples/submit_batch_generate.sh | 2 +- matey/generate.py | 11 ++-- matey/models/avit.py | 11 ++-- matey/models/basemodel.py | 9 +++- matey/models/svit.py | 11 ++-- matey/models/turbt.py | 11 ++-- matey/models/vit.py | 11 ++-- matey/train.py | 57 +++++++++------------ 12 files changed, 77 insertions(+), 98 deletions(-) diff --git a/examples/config/Demo_MW_diffusion_TT.yaml b/examples/config/Demo_MW_diffusion_TT.yaml index 00ca247..ec52378 100644 --- a/examples/config/Demo_MW_diffusion_TT.yaml +++ b/examples/config/Demo_MW_diffusion_TT.yaml @@ -44,12 +44,13 @@ basic_config: &basic_config 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: !!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) + 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]] diff --git a/examples/config/Demo_MW_diffusion_avit.yaml b/examples/config/Demo_MW_diffusion_avit.yaml index 1b9d5d6..b50e150 100644 --- a/examples/config/Demo_MW_diffusion_avit.yaml +++ b/examples/config/Demo_MW_diffusion_avit.yaml @@ -44,12 +44,13 @@ basic_config: &basic_config 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: !!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) + 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]] diff --git a/examples/config/Demo_MW_diffusion_svit.yaml b/examples/config/Demo_MW_diffusion_svit.yaml index 4896658..7c3a59c 100644 --- a/examples/config/Demo_MW_diffusion_svit.yaml +++ b/examples/config/Demo_MW_diffusion_svit.yaml @@ -44,12 +44,13 @@ basic_config: &basic_config 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: !!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) + 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]] diff --git a/examples/config/Demo_MW_diffusion_vit.yaml b/examples/config/Demo_MW_diffusion_vit.yaml index f99c8e7..736bf58 100644 --- a/examples/config/Demo_MW_diffusion_vit.yaml +++ b/examples/config/Demo_MW_diffusion_vit.yaml @@ -44,12 +44,13 @@ basic_config: &basic_config 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: !!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) + 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]] diff --git a/examples/submit_batch_generate.sh b/examples/submit_batch_generate.sh index dcb2e2a..4dbe598 100644 --- a/examples/submit_batch_generate.sh +++ b/examples/submit_batch_generate.sh @@ -16,7 +16,7 @@ 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 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/" diff --git a/matey/generate.py b/matey/generate.py index ff47241..e11b8ae 100644 --- a/matey/generate.py +++ b/matey/generate.py @@ -35,7 +35,8 @@ def __init__(self, params, global_rank, local_rank, device): self.epoch = 0 self.mp_type = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.half - self.cond_diffusion = getattr(self.params, "cond_diffusion", False) + 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 @@ -96,10 +97,10 @@ def initialize_data(self): self.single_print("self.train_data_loader:", len(self.train_data_loader), "valid_data_loader:", len(self.valid_data_loader)) def initialize_model(self): - if getattr(self.params, "diffusion", False): + if self.diffusion_config.get("diffusion", False): self.model = build_diffusion_model(self.params).to(self.device) else: - raise NotImplementedError("Only diffusion model generation is implemented currently. Please set params.diffusion to True.") + raise NotImplementedError("Only diffusion model generation is implemented currently. Please set params.diffusion_config.diffusion to True.") if dist.is_initialized() and self.params.use_ddp: @@ -254,7 +255,7 @@ def generate(self, seed=None, num_samples=1, batch_list=[0]): isgraph=isgraph, field_labels_out= field_labels ) - if getattr(self.params, "cond_diffusion", False): + if self.cond_diffusion: opts.diffusion_cond = rearrange(inp.to(self.device), 't b c d h w -> b t c d h w') # conditioning on input history # Euler step. # denoised = net(x_hat, t_hat, class_labels).to(torch.float64) @@ -283,7 +284,7 @@ def generate(self, seed=None, num_samples=1, batch_list=[0]): isgraph=isgraph, field_labels_out= field_labels ) - if getattr(self.params, "cond_diffusion", False): + if self.cond_diffusion: opts.diffusion_cond = rearrange(inp.to(self.device), 't b c d h w -> b t c d h w') # conditioning on input history # denoised = net(x_next, t_next, class_labels).to(torch.float64) # denoised = net(x_next, t_next, None) diff --git a/matey/models/avit.py b/matey/models/avit.py index b6f15a5..41bc164 100644 --- a/matey/models/avit.py +++ b/matey/models/avit.py @@ -38,11 +38,7 @@ def build_avit(params): n_steps=params.n_steps, bias_type=params.bias_type, hierarchical=getattr(params, 'hierarchical', None), - diffusion=getattr(params, 'diffusion', False), - model_channels=getattr(params, 'model_channels', 128), - channel_mult_emb=getattr(params, 'channel_mult_emb', 4), - embedding_type=getattr(params, 'embedding_type', 'positional'), - channel_mult_noise=getattr(params, 'channel_mult_noise', 1), + diffusion_config=getattr(params, 'diffusion_config', None), ) return model @@ -59,11 +55,10 @@ class AViT(BaseModel): """ 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, - diffusion=False, model_channels=128, channel_mult_emb=4, embedding_type='positional', channel_mult_noise=1): + 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, - diffusion=diffusion, model_channels=model_channels, channel_mult_emb=channel_mult_emb, - embedding_type=embedding_type, channel_mult_noise=channel_mult_noise) + diffusion_config=diffusion_config) self.drop_path = drop_path self.dp = np.linspace(0, drop_path, processor_blocks) diff --git a/matey/models/basemodel.py b/matey/models/basemodel.py index 803f6fb..d9cfa9e 100644 --- a/matey/models/basemodel.py +++ b/matey/models/basemodel.py @@ -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, - diffusion=False, model_channels=128, channel_mult_emb=4, embedding_type='positional', channel_mult_noise=1): + 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,8 +91,13 @@ 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.diffusion = diffusion + _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') diff --git a/matey/models/svit.py b/matey/models/svit.py index d0569d1..9079c00 100644 --- a/matey/models/svit.py +++ b/matey/models/svit.py @@ -39,11 +39,7 @@ def build_svit(params): bias_type=params.bias_type, replace_patch=getattr(params, 'replace_patch', True), hierarchical=getattr(params, 'hierarchical', None), - diffusion=getattr(params, 'diffusion', False), - model_channels=getattr(params, 'model_channels', 128), - channel_mult_emb=getattr(params, 'channel_mult_emb', 4), - embedding_type=getattr(params, 'embedding_type', 'positional'), - channel_mult_noise=getattr(params, 'channel_mult_noise', 1), + diffusion_config=getattr(params, 'diffusion_config', None), ) return model @@ -59,11 +55,10 @@ class sViT_all2all(BaseModel): """ 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, 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=False, model_channels=128, channel_mult_emb=4, embedding_type='positional', channel_mult_noise=1): + 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, - diffusion=diffusion, model_channels=model_channels, channel_mult_emb=channel_mult_emb, - embedding_type=embedding_type, channel_mult_noise=channel_mult_noise) + diffusion_config=diffusion_config) self.drop_path = drop_path self.dp = np.linspace(0, drop_path, processor_blocks) diff --git a/matey/models/turbt.py b/matey/models/turbt.py index 3898b76..06c327d 100644 --- a/matey/models/turbt.py +++ b/matey/models/turbt.py @@ -43,11 +43,7 @@ def build_turbt(params): replace_patch=getattr(params, 'replace_patch', True), hierarchical=getattr(params, 'hierarchical', None), notransposed=getattr(params, 'notransposed', False), - diffusion=getattr(params, 'diffusion', False), - model_channels=getattr(params, 'model_channels', 128), - channel_mult_emb=getattr(params, 'channel_mult_emb', 4), - embedding_type=getattr(params, 'embedding_type', 'positional'), - channel_mult_noise=getattr(params, 'channel_mult_noise', 1), + diffusion_config=getattr(params, 'diffusion_config', None), ) return model @@ -65,11 +61,10 @@ class TurbT(BaseModel): """ def __init__(self, tokenizer_heads=None, embed_dim=768, num_heads=12, processor_blocks=8, n_states=6, 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=False, model_channels=128, channel_mult_emb=4, embedding_type='positional', channel_mult_noise=1): + 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, - diffusion=diffusion, model_channels=model_channels, channel_mult_emb=channel_mult_emb, - embedding_type=embedding_type, channel_mult_noise=channel_mult_noise) + diffusion_config=diffusion_config) self.drop_path = drop_path self.dp = np.linspace(0, drop_path, processor_blocks) self.module_blocks=nn.ModuleDict({}) diff --git a/matey/models/vit.py b/matey/models/vit.py index 9bee6e8..659f055 100644 --- a/matey/models/vit.py +++ b/matey/models/vit.py @@ -38,11 +38,7 @@ def build_vit(params): replace_patch=getattr(params, 'replace_patch', True), hierarchical=getattr(params, 'hierarchical', None), use_linear=getattr(params, 'use_linear', False), - diffusion=getattr(params, 'diffusion', False), - model_channels=getattr(params, 'model_channels', 128), - channel_mult_emb=getattr(params, 'channel_mult_emb', 4), - embedding_type=getattr(params, 'embedding_type', 'positional'), - channel_mult_noise=getattr(params, 'channel_mult_noise', 1), + diffusion_config=getattr(params, 'diffusion_config', None), ) return model @@ -60,11 +56,10 @@ class ViT_all2all(BaseModel): """ def __init__(self, tokenizer_heads=None, embed_dim=768, 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", replace_patch=True, SR_ratio=[1,1,1], hierarchical=None, use_linear=False, - diffusion=False, model_channels=128, channel_mult_emb=4, embedding_type='positional', channel_mult_noise=1): + 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, use_linear=use_linear, - diffusion=diffusion, model_channels=model_channels, channel_mult_emb=channel_mult_emb, - embedding_type=embedding_type, channel_mult_noise=channel_mult_noise) + diffusion_config=diffusion_config) self.drop_path = drop_path self.dp = np.linspace(0, drop_path, processor_blocks) self.blocks = nn.ModuleList([SpaceTimeBlock_all2all(embed_dim, num_heads,drop_path=self.dp[i]) diff --git a/matey/train.py b/matey/train.py index f14bd29..f30046e 100644 --- a/matey/train.py +++ b/matey/train.py @@ -41,6 +41,7 @@ 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)) @@ -94,7 +95,7 @@ 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 getattr(self.params, "diffusion", False): + if self.diffusion_config.get("diffusion", False): self.diffusion_loss = EDMLoss() self.initialize_model() @@ -183,7 +184,7 @@ def initialize_data(self): self.val_sampler.set_epoch(0) def initialize_model(self): - if getattr(self.params, "diffusion", False): + 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) @@ -535,7 +536,7 @@ 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 getattr(self.params, "diffusion", False): + 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: @@ -553,17 +554,12 @@ def train_one_epoch(self): data_time = 0 data_start = self.timer.get_time() self.model.train() - if getattr(self.params, "diffusion", False): - logs = {'train_EDMloss': torch.zeros(1).to(self.device), - '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)} - else: - logs = {'train_rmse': torch.zeros(1).to(self.device), - 'train_nrmse': torch.zeros(1).to(self.device), + 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)} + 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)) @@ -642,21 +638,20 @@ def train_one_epoch(self): field_labels_out= field_labels_out ) with record_function_opt("model forward", enabled=self.profiling): - if getattr(self.params, "diffusion", False): - if getattr(self.params, "cond_diffusion", False): + 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') - 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) else: assert torch.all(inp[0] == tar), "For unconditioned diffusion model, input should be the target" - output, loss_field = self.model_forward(inp, field_labels, bcs, opts) + 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 getattr(self.params, "diffusion", False): + 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: @@ -737,17 +732,12 @@ def train_one_epoch(self): def validate_one_epoch(self, full=False, cutoff_skip=False): self.model.eval() self.single_print('STARTING VALIDATION!!!') - if getattr(self.params, "diffusion", False): - logs = {'valid_EDMloss': torch.zeros(1).to(self.device), - 'valid_rmse': torch.zeros(1).to(self.device), - 'valid_nrmse': torch.zeros(1).to(self.device), - 'valid_l1': torch.zeros(1).to(self.device), - 'valid_ssim': torch.zeros(1).to(self.device)} - else: - logs = {'valid_rmse': torch.zeros(1).to(self.device), - 'valid_nrmse': torch.zeros(1).to(self.device), - 'valid_l1': torch.zeros(1).to(self.device), - 'valid_ssim': torch.zeros(1).to(self.device)} + logs = {'valid_rmse': torch.zeros(1).to(self.device), + '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} @@ -828,14 +818,13 @@ def validate_one_epoch(self, full=False, cutoff_skip=False): isgraph=isgraph, field_labels_out= field_labels_out ) - if getattr(self.params, "diffusion", False): - if getattr(self.params, "cond_diffusion", False): + 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') - 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) else: assert torch.all(inp[0] == tar), "For unconditioned diffusion model, input should be the target" - output, loss_field = self.model_forward(inp, field_labels, bcs, opts) + 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) From 8f26d0951a9a7f3d8c1b8c382c656efbe34361bc Mon Sep 17 00:00:00 2001 From: Paul Laiu Date: Thu, 21 May 2026 18:24:40 -0400 Subject: [PATCH 10/12] Renamed diffusion model specific network blocks; moved diffusion noise embedding computation to basemodel.py --- matey/models/avit.py | 7 +------ matey/models/basemodel.py | 23 +++++++++++++++-------- matey/models/svit.py | 8 ++------ matey/models/time_modules.py | 6 +++--- matey/models/turbt.py | 8 ++------ matey/models/vit.py | 6 +----- 6 files changed, 24 insertions(+), 34 deletions(-) diff --git a/matey/models/avit.py b/matey/models/avit.py index 41bc164..3fe909f 100644 --- a/matey/models/avit.py +++ b/matey/models/avit.py @@ -7,7 +7,6 @@ import torch.nn.functional as F import numpy as np from einops import rearrange, repeat -from torch.nn.functional import silu from .spacetime_modules import SpaceTimeBlock from .basemodel import BaseModel from ..data_utils.shared_utils import normalize_spatiotemporal_persample, get_top_variance_patchids @@ -178,11 +177,7 @@ def forward(self, x, state_labels, bcs, opts: ForwardOptionsBase, train_opts: Op T, _, _, D, _, _ = x.shape if self.diffusion: - emb = self.map_noise(sigma) - emb = emb.reshape(emb.shape[0], 2, -1).flip(1).reshape(*emb.shape) - emb = silu(self.map_layer0(emb)) - emb = silu(self.map_layer1(emb)) - emb = self.affine[str(imod)](emb) # (B, embed_dim) + 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: diff --git a/matey/models/basemodel.py b/matey/models/basemodel.py index d9cfa9e..6d90c18 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, FourierEmbedding, PositionalEmbedding, Linear +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 @@ -102,19 +102,19 @@ def __init__(self, tokenizer_heads, n_states=6, n_states_out=None, n_states_cond noise_channels = model_channels * channel_mult_noise init = dict(init_mode='xavier_uniform') - self.map_noise = (PositionalEmbedding(num_channels=noise_channels, endpoint=True) + self.map_noise = (PositionalEmbedding_EDM(num_channels=noise_channels, endpoint=True) if embedding_type == 'positional' - else FourierEmbedding(num_channels=noise_channels)) - self.map_layer0 = Linear(in_features=noise_channels, out_features=emb_channels, **init) - self.map_layer1 = Linear(in_features=emb_channels, out_features=emb_channels, **init) + 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(in_features=emb_channels, out_features=embed_dim, **init) + 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(in_features=embed_dim, out_features=embed_dim, **init) + self.diffusion_cond_proj[str(ilevel)] = Linear_EDM(in_features=embed_dim, out_features=embed_dim, **init) self.embed_dim=embed_dim @@ -703,9 +703,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/svit.py b/matey/models/svit.py index 9079c00..78382ec 100644 --- a/matey/models/svit.py +++ b/matey/models/svit.py @@ -6,7 +6,7 @@ import torch.nn as nn import numpy as np from einops import rearrange -from torch.nn.functional import silu + 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 @@ -147,11 +147,7 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase, train_opts: field_labels_out = state_labels if self.diffusion: - emb = self.map_noise(sigma) - emb = emb.reshape(emb.shape[0], 2, -1).flip(1).reshape(*emb.shape) - emb = silu(self.map_layer0(emb)) - emb = silu(self.map_layer1(emb)) - emb = self.affine[str(imod)](emb) # (B, embed_dim) + 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') diff --git a/matey/models/time_modules.py b/matey/models/time_modules.py index 52c74b4..011769b 100644 --- a/matey/models/time_modules.py +++ b/matey/models/time_modules.py @@ -19,7 +19,7 @@ def forward(self, 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(torch.nn.Module): +class PositionalEmbedding_EDM(torch.nn.Module): def __init__(self, num_channels, max_positions=10000, endpoint=False): super().__init__() self.num_channels = num_channels @@ -34,7 +34,7 @@ def forward(self, x): x = torch.cat([x.cos(), x.sin()], dim=1) return x -class FourierEmbedding(torch.nn.Module): +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) @@ -51,7 +51,7 @@ def weight_init(shape, mode, fan_in, fan_out): if mode == 'kaiming_normal': return np.sqrt(1 / fan_in) * torch.randn(*shape) raise ValueError(f'Invalid init mode "{mode}"') -class Linear(torch.nn.Module): +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 diff --git a/matey/models/turbt.py b/matey/models/turbt.py index 06c327d..56d428c 100644 --- a/matey/models/turbt.py +++ b/matey/models/turbt.py @@ -12,7 +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 -from torch.nn.functional import silu + import sys, copy from operator import mul from functools import reduce @@ -239,11 +239,7 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase): field_labels_out = state_labels if self.diffusion: - emb = self.map_noise(sigma) - emb = emb.reshape(emb.shape[0], 2, -1).flip(1).reshape(*emb.shape) # swap sin/cos - emb = silu(self.map_layer0(emb)) - emb = silu(self.map_layer1(emb)) - emb = self.affine[str(imod)](emb) # (B, embed_dim) + 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') diff --git a/matey/models/vit.py b/matey/models/vit.py index 659f055..f0f6653 100644 --- a/matey/models/vit.py +++ b/matey/models/vit.py @@ -145,11 +145,7 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase, train_opts: field_labels_out = state_labels if self.diffusion: - emb = self.map_noise(sigma) - emb = emb.reshape(emb.shape[0], 2, -1).flip(1).reshape(*emb.shape) - emb = silu(self.map_layer0(emb)) - emb = silu(self.map_layer1(emb)) - emb = self.affine[str(imod)](emb) # (B, embed_dim) + 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') From 760af1c0e5135c81c8025e7639cf1efd58bcdb91 Mon Sep 17 00:00:00 2001 From: Paul Laiu Date: Thu, 21 May 2026 18:25:52 -0400 Subject: [PATCH 11/12] batch sample generation --- matey/generate.py | 169 +++++++++++++++++++++------------------------- 1 file changed, 77 insertions(+), 92 deletions(-) diff --git a/matey/generate.py b/matey/generate.py index e11b8ae..9c258af 100644 --- a/matey/generate.py +++ b/matey/generate.py @@ -203,110 +203,95 @@ def generate(self, seed=None, num_samples=1, batch_list=[0]): 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")}') - for sample_idx in range(num_samples): - - - init_inp = torch.randn(inp.shape).to(self.device) - - - self.model.eval() - - with torch.no_grad(): - - - sigma_min = 0.002 - sigma_max = 80 - n_diffusion_steps = 18 - rho=7 - S_churn=0 - S_min=0 - S_max=float('inf') - S_noise=1 - - step_indices = torch.arange(n_diffusion_steps, device=self.device) - - t_steps = (sigma_max ** (1 / rho) + step_indices / (n_diffusion_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho - t_steps = torch.cat([t_steps, torch.zeros_like(t_steps[:1])]) # t_N = 0 - - # tar = tar.to(self.device) - # init_inp = rearrange(init_inp.to(self.device), 'b t c d h w -> t b c d h w') - - - x_next = init_inp * t_steps[0] - - for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])): # 0, ..., N-1 - x_cur = x_next - - # Increase noise temporarily. - gamma = min(S_churn / n_diffusion_steps, np.sqrt(2) - 1) if S_min <= t_cur <= S_max else 0 - t_hat = t_cur + gamma * t_cur - x_hat = x_cur + (t_hat ** 2 - t_cur ** 2).sqrt() * S_noise * torch.randn_like(x_cur) - - - opts = ForwardOptionsBase( - imod=imod, - imod_bottom=imod_bottom , + self.model.eval() + + with torch.no_grad(): + sigma_min = 0.002 + sigma_max = 80 + n_diffusion_steps = 18 + rho=7 + S_churn=0 + S_min=0 + S_max=float('inf') + S_noise=1 + + step_indices = torch.arange(n_diffusion_steps, device=self.device) + t_steps = (sigma_max ** (1 / rho) + step_indices / (n_diffusion_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho + t_steps = torch.cat([t_steps, torch.zeros_like(t_steps[:1])]) # t_N = 0 + + # 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) + + 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:])): # 0, ..., N-1 + x_cur = x_next + + # Increase noise temporarily. + gamma = min(S_churn / n_diffusion_steps, np.sqrt(2) - 1) if S_min <= t_cur <= S_max else 0 + t_hat = t_cur + gamma * t_cur + x_hat = x_cur + (t_hat ** 2 - t_cur ** 2).sqrt() * S_noise * torch.randn_like(x_cur) + + opts = ForwardOptionsBase( + imod=imod, + imod_bottom=imod_bottom, tkhead_name=tkhead_name, sequence_parallel_group=seq_group, - leadtime=leadtime, + leadtime=leadtime_b, blockdict=copy.deepcopy(blockdict), - cond_dict=copy.deepcopy(cond_dict), - cond_input=cond_input, + cond_dict=copy.deepcopy(cond_dict_b), + cond_input=cond_input_b, isgraph=isgraph, - field_labels_out= field_labels - ) - if self.cond_diffusion: - opts.diffusion_cond = rearrange(inp.to(self.device), 't b c d h w -> b t c d h w') # conditioning on input history - # Euler step. - # denoised = net(x_hat, t_hat, class_labels).to(torch.float64) - # denoised = net(x_hat, t_hat, None) - - # denoised = self.inference(x_hat, t_hat, field_labels, bcs, opts) - denoised = self.model(x_hat, t_hat.repeat(x_hat.shape[1]), field_labels, bcs, opts) - - # print(f"denoised shape: {denoised.shape}, x_hat shape: {x_hat.shape}, t_hat shape: {t_hat.shape}") - d_cur = (x_hat - denoised) / t_hat - x_next = x_hat + (t_next - t_hat) * d_cur - - # print(f"x_next shape: {x_next.shape}, d_cur shape: {d_cur.shape}, t_next shape: {t_next.shape}, t_hat shape: {t_hat.shape}") - # Apply 2nd order correction. - if i < n_diffusion_steps - 1: - - opts = ForwardOptionsBase( - imod=imod, - imod_bottom=imod_bottom , + field_labels_out=field_labels_b + ) + if self.cond_diffusion: + opts.diffusion_cond = rearrange(inp_b, 't b c d h w -> b t c d h w') + + # Euler step. + denoised = self.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 + + # Apply 2nd order correction. + if i < n_diffusion_steps - 1: + opts = ForwardOptionsBase( + imod=imod, + imod_bottom=imod_bottom, tkhead_name=tkhead_name, sequence_parallel_group=seq_group, - leadtime=leadtime, + leadtime=leadtime_b, blockdict=copy.deepcopy(blockdict), - cond_dict=copy.deepcopy(cond_dict), - cond_input=cond_input, + cond_dict=copy.deepcopy(cond_dict_b), + cond_input=cond_input_b, isgraph=isgraph, - field_labels_out= field_labels - ) - if self.cond_diffusion: - opts.diffusion_cond = rearrange(inp.to(self.device), 't b c d h w -> b t c d h w') # conditioning on input history - # denoised = net(x_next, t_next, class_labels).to(torch.float64) - # denoised = net(x_next, t_next, None) - # denoised = self.inference(x_next, t_next, field_labels, bcs, opts) - denoised = self.model(x_next, t_next.repeat(x_next.shape[1]), field_labels, bcs, opts) - - d_prime = (x_next - denoised) / t_next - x_next = x_hat + (t_next - t_hat) * (0.5 * d_cur + 0.5 * d_prime) + field_labels_out=field_labels_b + ) + if self.cond_diffusion: + opts.diffusion_cond = rearrange(inp_b, 't b c d h w -> b t c d h w') - torch.save(x_next.cpu(), os.path.join(self.output_dir, f'generation_step_{i}_batch_{batch_idx}_sample{sample_idx}.pt')) + denoised = self.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) - output = x_next - - ###full resolution### - # residuals = output - tar - # torch.save(output.cpu(), os.path.join('Demo_Diffusion_CIFAR10_finemodel/basic_config/demo_diffusion/training_checkpoints/', 'generation_output.pt')) - # self.single_print(f'Generation output saved to {"Demo_Diffusion_CIFAR10_finemodel/basic_config/demo_diffusion/training_checkpoints/generation_output.pt"}') - torch.save(output.cpu(), os.path.join(self.output_dir, f'generation_output_batch_{batch_idx}_sample{sample_idx}.pt')) - self.single_print(f'Generation output saved to {os.path.join(self.output_dir, f"generation_output_batch_{batch_idx}_sample{sample_idx}.pt")}') + # Save all samples for this step: [T, B*S, C, D, H, W] -> [S, T, B, C, D, H, W] + x_next_grouped = rearrange(x_next, 't (b s) c d h w -> s t b c d h w', s=num_samples) + torch.save(x_next_grouped.cpu(), os.path.join(self.output_dir, f'generation_step_{i}_batch_{batch_idx}.pt')) + # Save final output: [T, B*S, C, D, H, W] -> [S, T, B, C, D, H, W] + output = rearrange(x_next, 't (b s) c d h w -> s t b c d h w', s=num_samples) + 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() + torch.cuda.empty_cache() From efa14d5de81a2d59f7088b6aaa8de6f9462a1c0c Mon Sep 17 00:00:00 2001 From: Paul Laiu Date: Thu, 21 May 2026 18:53:20 -0400 Subject: [PATCH 12/12] make a sampler class in generate.py for the EDM sampler --- matey/generate.py | 277 ++++++++++++++++++++++------------------------ 1 file changed, 132 insertions(+), 145 deletions(-) diff --git a/matey/generate.py b/matey/generate.py index 9c258af..59450f9 100644 --- a/matey/generate.py +++ b/matey/generate.py @@ -21,6 +21,80 @@ 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 @@ -30,7 +104,8 @@ def __init__(self, params, global_rank, local_rank, device): 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.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 @@ -206,19 +281,6 @@ def generate(self, seed=None, num_samples=1, batch_list=[0]): self.model.eval() with torch.no_grad(): - sigma_min = 0.002 - sigma_max = 80 - n_diffusion_steps = 18 - rho=7 - S_churn=0 - S_min=0 - S_max=float('inf') - S_noise=1 - - step_indices = torch.arange(n_diffusion_steps, device=self.device) - t_steps = (sigma_max ** (1 / rho) + step_indices / (n_diffusion_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho - t_steps = torch.cat([t_steps, torch.zeros_like(t_steps[:1])]) # t_N = 0 - # 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) @@ -231,63 +293,25 @@ def generate(self, seed=None, num_samples=1, batch_list=[0]): 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) - 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:])): # 0, ..., N-1 - x_cur = x_next - - # Increase noise temporarily. - gamma = min(S_churn / n_diffusion_steps, np.sqrt(2) - 1) if S_min <= t_cur <= S_max else 0 - t_hat = t_cur + gamma * t_cur - x_hat = x_cur + (t_hat ** 2 - t_cur ** 2).sqrt() * S_noise * torch.randn_like(x_cur) - - opts = ForwardOptionsBase( - imod=imod, - imod_bottom=imod_bottom, - tkhead_name=tkhead_name, - sequence_parallel_group=seq_group, - leadtime=leadtime_b, - blockdict=copy.deepcopy(blockdict), - cond_dict=copy.deepcopy(cond_dict_b), - cond_input=cond_input_b, - isgraph=isgraph, - field_labels_out=field_labels_b - ) - if self.cond_diffusion: - opts.diffusion_cond = rearrange(inp_b, 't b c d h w -> b t c d h w') - - # Euler step. - denoised = self.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 - - # Apply 2nd order correction. - if i < n_diffusion_steps - 1: - opts = ForwardOptionsBase( - imod=imod, - imod_bottom=imod_bottom, - tkhead_name=tkhead_name, - sequence_parallel_group=seq_group, - leadtime=leadtime_b, - blockdict=copy.deepcopy(blockdict), - cond_dict=copy.deepcopy(cond_dict_b), - cond_input=cond_input_b, - isgraph=isgraph, - field_labels_out=field_labels_b - ) - if self.cond_diffusion: - opts.diffusion_cond = rearrange(inp_b, 't b c d h w -> b t c d h w') - - denoised = self.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) - - # Save all samples for this step: [T, B*S, C, D, H, W] -> [S, T, B, C, D, H, W] - x_next_grouped = rearrange(x_next, 't (b s) c d h w -> s t b c d h w', s=num_samples) - torch.save(x_next_grouped.cpu(), os.path.join(self.output_dir, f'generation_step_{i}_batch_{batch_idx}.pt')) - - # Save final output: [T, B*S, C, D, H, W] -> [S, T, B, C, D, H, W] - output = rearrange(x_next, 't (b s) c d h w -> s t b c d h w', s=num_samples) + 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")}') @@ -353,80 +377,43 @@ def autoregressive_generate(self, seed=None, num_samples=1, num_steps=10): inp_cur = inp.clone() # conditioning that advances each autoregressive step for step_idx in range(num_steps): - samples = [] # collect num_samples independent draws at this step - - for sample_idx in range(num_samples): - with torch.no_grad(): - sigma_min = 0.002 - sigma_max = 80 - n_diffusion_steps = 18 - rho = 7 - S_churn = 0 - S_min = 0 - S_max = float('inf') - S_noise = 1 - - step_indices = torch.arange(n_diffusion_steps, device=self.device) - t_steps = (sigma_max ** (1 / rho) + step_indices / (n_diffusion_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho - t_steps = torch.cat([t_steps, torch.zeros_like(t_steps[:1])]) # t_N = 0 - - x_next = torch.randn(inp_cur.shape, device=self.device) * t_steps[0] - - for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])): - x_cur = x_next - - # Increase noise temporarily. - gamma = min(S_churn / n_diffusion_steps, np.sqrt(2) - 1) if S_min <= t_cur <= S_max else 0 - t_hat = t_cur + gamma * t_cur - x_hat = x_cur + (t_hat ** 2 - t_cur ** 2).sqrt() * S_noise * torch.randn_like(x_cur) - - opts = ForwardOptionsBase( - imod=imod, - imod_bottom=imod_bottom, - tkhead_name=tkhead_name, - sequence_parallel_group=seq_group, - leadtime=leadtime, - blockdict=copy.deepcopy(blockdict), - cond_dict=copy.deepcopy(cond_dict), - cond_input=cond_input, - isgraph=isgraph, - field_labels_out=field_labels - ) - if self.cond_diffusion: - opts.diffusion_cond = rearrange(inp_cur, 't b c d h w -> b t c d h w') - - # Euler step. - denoised = self.model(x_hat, t_hat.repeat(x_hat.shape[1]), field_labels, bcs, opts) - d_cur = (x_hat - denoised) / t_hat - x_next = x_hat + (t_next - t_hat) * d_cur - - # Apply 2nd order correction. - if i < n_diffusion_steps - 1: - opts = ForwardOptionsBase( - imod=imod, - imod_bottom=imod_bottom, - tkhead_name=tkhead_name, - sequence_parallel_group=seq_group, - leadtime=leadtime, - blockdict=copy.deepcopy(blockdict), - cond_dict=copy.deepcopy(cond_dict), - cond_input=cond_input, - isgraph=isgraph, - field_labels_out=field_labels - ) - if self.cond_diffusion: - opts.diffusion_cond = rearrange(inp_cur, 't b c d h w -> b t c d h w') - - denoised = self.model(x_next, t_next.repeat(x_next.shape[1]), field_labels, bcs, opts) - d_prime = (x_next - denoised) / t_next - x_next = x_hat + (t_next - t_hat) * (0.5 * d_cur + 0.5 * d_prime) - - samples.append(x_next) - torch.save(x_next.cpu(), os.path.join(self.output_dir, f'generation_output_step_{step_idx}_sample{sample_idx}.pt')) - self.single_print(f'Generation output saved to {os.path.join(self.output_dir, f"generation_output_step_{step_idx}_sample{sample_idx}.pt")}') - - # Mean of all samples becomes the conditioning for the next autoregressive step - inp_cur = torch.stack(samples, dim=0).mean(dim=0) + 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")}')