ldcast-dmi/example_config.yaml at master · dmidk/ldcast-dmi · GitHub

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
schema_version: v0.1.0

general:
  seed: 24
  max_nproc: 4
  model_dir: &model_dir /vf/ldcast/test_models/

datasets:
  radar_composite: &radar_composite
    path: /dmidata/projects/nwp/radar/tmp_ldcast_out/zarr_patch/test_radar_2017_07.zarr
  radar_predict: &radar_predict
    path: /vf/ldcast/demo/20210622

preprocessing:
  sampling:
    sampler: EqualFrequencySampler
    sample_shape: [2, 2]
    timestep: PT10M
    num_timesteps: 2
    time_range_sampling: [-PT10M, PT20M]
    overwrite: False
    inputs:
      dataset: *radar_composite
    outputs: &radar_sampled
      path: /dmidata/projects/nwp/radar/tmp_ldcast_out/sampled_zarr/sampled_test_radar_2017_07.zarr

models:
  autoenc:
    batch_size: 4
    batches_per_epoch: 2 # Should be a multiple of the number of GPUs
    max_epochs: 2
    chpt_path: null
    inputs:
      datasets:
        sampled: *radar_sampled
        target: *radar_composite
        predictions:
          radar_composite: *radar_composite
    outputs:
      autoenc_weights: &autoenc_weights
        filename: autoenc-test.pt
        directory: *model_dir

  genforecast:
    batch_size: 1
    batches_per_epoch: 2 # Should be a multiple of the number of GPUs
    max_epochs: 2
    chpt_path: null
    num_future_timesteps: 3
    num_past_timesteps: 3
    use_obs: True
    use_nwp: False
    lr: 1e-4
    inputs:
      datasets:
        sampled: *radar_sampled
        target: *radar_composite
        predictions:
          radar_composite: *radar_composite
      autoenc_weights: *autoenc_weights
    outputs:
      genforecast_weights: &genforecast_weights
        filename: genforecast-test.pt
        directory: *model_dir

  forecast:
    num_diffusion_iters: 2
    timestep: PT5M
    past_timesteps: 3
    future_timesteps: 3
    autoenc_time_ratio: 1
    crop_box: [[128, 480], [160, 608]]
    inputs:
      dataset: *radar_predict
      autoenc_weights: *autoenc_weights
      genforecast_weights: *genforecast_weights
    outputs:
      path: /vf/ldcast/demo/forecast/test_forecast.zarr