name: AE
workspace: gaoch
project: diffusion_main

hyperparameters:
# batch要求必须放在这里
  global_batch_size: 16
  model_params:
    block_expansion: 64
    max_features: 512
    num_down_blocks: 2
    num_up_blocks: 2
    num_bottleneck_blocks: 6
    scales: [1, 0.5, 0.25, 0.125]
    skips: True
    loss_weights:
      perceptual: [10, 10, 10, 10, 10]
      equivariance_shift: 10
      equivariance_affine: 10
    flow_predictor_params:
      num_channels: 4
      num_blocks: 2
      max_features: 512
      block_expansion: 64
      flow_type: z

  train_params:
    epochs: 500
    step_size: 100
    adam_betas: [0.9, 0.99]
    only_use_flow: True
    # epoch_milestones: [800, 1000]
    max_epochs: 100
    num_repeats: 100
    epoch_milestones: [60, 90]
    lr: 2.0e-4
    dataloader_workers: 6
    print_freq: 10
    save_img_freq: 100
    update_ckpt_freq: 5000
    scales: [1, 0.5, 0.25, 0.125]
    transform_params:
      sigma_affine: 0.05
      sigma_tps: 0.005
      points_tps: 5
    loss_weights:
      perceptual: [10, 10, 10, 10, 10]
      equivariance_shift: 10
      equivariance_affine: 10
  data_params:
    # data_path: 2D_CFD_Rand_M0.1_Eta0.01_Zeta0.01_periodic_128_Train
    # data_path: 2D_CFD_Rand_M0.1_Eta0.1_Zeta0.1_periodic_128_Train
    # data_path: 2D_CFD_Rand_M0.1_Eta1e-08_Zeta1e-08_periodic_512_Train

    # data_path: 2D_CFD_Rand_M1.0_Eta0.01_Zeta0.01_periodic_128_Train
    # data_path: 2D_CFD_Rand_M1.0_Eta0.1_Zeta0.1_periodic_128_Train
    # data_path: 2D_CFD_Rand_M1.0_Eta1e-08_Zeta1e-08_periodic_512_Train

    # data_path: 2D_CFD_Turb_M0.1_Eta1e-08_Zeta1e-08_periodic_512_Train

    # data_path: 2D_CFD_Turb_M1.0_Eta1e-08_Zeta1e-08_periodic_512_Train
    # data_path: PDEbench/2D/Incom/2D_diff-react_NA_NA.h5
    data_path: FluidDataset/doubleshock2npy
    download: false
    field_size: 256
    num_frames: 21
    num_channels: 2
    frame_shape: 256
    id_sampling: True
    pairs_list: null
    augmentation_params:
      flip_param:
        horizontal_flip: True
        time_flip: True
      jitter_param:
        brightness: 0.1
        contrast: 0.1
        saturation: 0.1
        hue: 0.1
  visualizer_params:
    kp_size: 2
    draw_border: True
    colormap: 'gist_rainbow'
    region_bg_color: [1, 1, 1]
  data_downloaded: false

searcher:
  name: single
  metric: validation_loss
  max_length:
      epochs: 250 #60,000 training images with batch size 64
  smaller_is_better: true

entrypoint: trainer_AE:AETrial
# entrypoint: python test.py

checkpoint_storage:
  access_key: det_master_shahe
  bucket: determined-checkpoint-storage
  endpoint_url: http://192.168.5.174:9000
  prefix: null
  save_experiment_best: 0
  save_trial_best: 1
  save_trial_latest: 1
  secret_key: k5uKFQIj8mobuugcAXbgVsqf7TsQstRjs2mSLwdt
  type: s3

environment:
  image:
    cpu: 192.168.5.174:5000/gaoch:v4
    cuda: 192.168.5.174:5000/gaoch:v4
    rocm: 192.168.5.174:5000/gaoch:v4
    # cpu: 100.64.0.3:5000/gaoch:v4
    # cuda: 100.64.0.3:5000/gaoch:v4
    # rocm: 100.64.0.3:5000/gaoch:v4

resources:
  devices: null
  resource_pool: RTX4090-X
  # resource_pool: RTXA6000
  slots_per_trial: 1
  weight: 5

min_validation_period:
   epochs: 5