File size: 2,088 Bytes
6021dd1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
from stldm.stldm import model_setup, guidance_scheduler, ae_setup, backbone_setup
from stldm.stldm_spatial import model_setup as spatial_setup
n2n_setup = {'2D': spatial_setup, '3D': model_setup}
STLDM_SEVIR = {
'model': "stldm",
'pre': None,
'post': None,
'vp_param': {
'shape_in': (13, 1, 128, 128),
'shape_out': (12, 1, 128, 128),
'hid_S': 32,
'hid_T': 512,
'N_S': 4,
'N_T': 8,
'groups': 8,
'last_activation': 'sigmoid',
},
'stldm_param': {
'in_ch': 32,
'chs_mult': [1,2,4,8],
'num_groups': 8,
'heads': 4,
'dim_head': 32,
'base_ch': 64,
'patch_size': 16
},
'param': {
'timesteps': 50,
'sampling_timesteps': 20,
'objective': 'pred_v'
}
}
STLDM_HKO = {
'model': "stldm",
'pre': None,
'post': None,
'vp_param': {
'shape_in': (5, 1, 128, 128),
'shape_out': (20, 1, 128, 128),
'hid_S': 32,
'hid_T': 512,
'N_S': 4,
'N_T': 8,
'groups': 8,
'last_activation': 'sigmoid',
},
'stldm_param': {
'in_ch': 32,
'chs_mult': [1,2,4,8],
'num_groups': 8,
'heads': 4,
'dim_head': 32,
'base_ch': 64,
'patch_size': 16
},
'param': {
'timesteps': 50,
'sampling_timesteps': 20,
'objective': 'pred_v'
}
}
STLDM_METEO = {
'model': "stldm",
'pre': None,
'post': None,
'vp_param': {
'shape_in': (5, 1, 128, 128),
'shape_out': (20, 1, 128, 128),
'hid_S': 32,
'hid_T': 512,
'N_S': 4,
'N_T': 8,
'groups': 8,
'last_activation': 'sigmoid',
},
'stldm_param': {
'in_ch': 32,
'chs_mult': [1,2,4,8],
'num_groups': 8,
'heads': 4,
'dim_head': 32,
'base_ch': 64,
'patch_size': 16
},
'param': {
'timesteps': 50,
'sampling_timesteps': 20,
'objective': 'pred_v'
}
} |