Loss of resolution in predictions

[robmarkcole]

The pixel regression outputs from the model lack resolution, I will here document efforts to resolve this at the suggestion of @APatrickJ

Initial - LARGE model

Left is truth, right is prediction

conv_layers_per_resolution at 1 down from 2

Some structure appears more clearly:

Base model

Now set transformer patch_size: 2 and input patch/crop size 16 -> created allenai/rslearn#484

Level of detail is improved:

[favyen2]

How many input images are you using btw?

[robmarkcole]

The total dataset is around 300 images (truth)

[gabrieltseng]

Hi @robmarkcole ,

Thanks for sharing this. Do you have more information about this (ideally the config used to train the model, if it was done using rslearn)?

Also - are there (different) models / recipes you have used which better preserve the resolution?

[robmarkcole]

@gabrieltseng I have been modifying this config - I've got best results with the large model:

model:
  class_path: rslearn.train.lightning_module.RslearnLightningModule
  init_args:
    model:
      class_path: rslearn.models.singletask.SingleTaskModel
      init_args:
        encoder:
          - class_path: rslearn.models.olmoearth_pretrain.model.OlmoEarth
            init_args:
              model_id: OLMOEARTH_V1_LARGE
              patch_size: 4
              use_legacy_timestamps: false
        decoder:
          - class_path: rslearn.models.unet.UNetDecoder
            init_args:
              # OLMOEARTH_V1_LARGE produces 1024 channels at factor=4.
              in_channels: [[4, 1024]]
              out_channels: 1
              conv_layers_per_resolution: 2
              num_channels: {8: 512, 4: 512, 2: 256, 1: 128}
          - class_path: rslearn.train.tasks.per_pixel_regression.PerPixelRegressionHead
            init_args:
              loss_mode: mse
    optimizer:
      class_path: rslearn.train.optimizer.AdamW
      init_args:
        lr: 0.00002
        weight_decay: 0.01
    scheduler:
      class_path: rslearn.train.scheduler.PlateauScheduler
      init_args:
        factor: 0.2
        patience: 2
        min_lr: 0
        cooldown: 10
data:
  class_path: rslearn.train.data_module.RslearnDataModule
  init_args:
    path: data-usa
    batch_size: 1
    num_workers: 4
    inputs:
      sentinel2_l2a:
        data_type: raster
        layers: ["sentinel2"]
        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12", "B01", "B09"]
        passthrough: true
        dtype: FLOAT32
        load_all_item_groups: true
        load_all_layers: true
      landsat:
        data_type: raster
        layers: ["landsat"]
        bands: ["B1", "B2", "B3", "B4", "B5", "B6", "B7", "B8", "B9", "B10", "B11"]
        passthrough: true
        dtype: FLOAT32
        load_all_item_groups: true
        load_all_layers: true
      targets:
        data_type: raster
        layers: ["yield_raster"]
        bands: ["yield"]
        dtype: FLOAT32
        is_target: true
    task:
      class_path: rslearn.train.tasks.per_pixel_regression.PerPixelRegressionTask
      init_args:
        scale_factor: 1.0
        metrics: ["mse", "r2"]
        nodata_value: -1
    default_config:
      crop_size: 32
      transforms:
        - class_path: rslearn.models.olmoearth_pretrain.norm.OlmoEarthNormalize
          init_args:
            band_names:
              sentinel2_l2a: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12", "B01", "B09"]
              landsat: ["B1", "B2", "B3", "B4", "B5", "B6", "B7", "B8", "B9", "B10", "B11"]
    train_config:
      groups: ["yield_fields"]
      tags:
        split: train
    val_config:
      groups: ["yield_fields"]
      tags:
        split: val
    test_config:
      groups: ["yield_fields"]
      tags:
        split: test
    predict_config:
      # Sliding-window inference over all crops of each window, stitched back to a
      # full-window raster by RslearnWriter below.
      load_all_crops: true
      crop_size: 32
      overlap_pixels: 8
      skip_targets: true
      output_layer_name_skip_inference_if_exists: yield_pred
      groups: ["yield_fields"]
      tags:
        split: test
trainer:
  max_epochs: 50
  callbacks:
    - class_path: lightning.pytorch.callbacks.LearningRateMonitor
      init_args:
        logging_interval: epoch
    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
      init_args:
        save_top_k: 1
        save_last: true
        # monitor: val_mse
        # mode: min
        monitor: val_r2
        mode: max
    - class_path: rslearn.train.prediction_writer.RslearnWriter
      init_args:
        path: data-usa
        output_layer: yield_pred
        merger:
          class_path: rslearn.train.prediction_writer.RasterMerger
          init_args:
            overlap_pixels: 8
            downsample_factor: 1
    - class_path: rslearn.train.callbacks.freeze_unfreeze.FreezeUnfreeze
      init_args:
        # Freeze the OlmoEarth encoder initially, then unfreeze for full fine-tuning.
        module_selector: ["model", "encoder", 0]
        unfreeze_at_epoch: 5
        unfreeze_lr_factor: 5
  logger:
    class_path: lightning.pytorch.loggers.WandbLogger
    init_args:
      project: olmoearth-crop-yield
      name: ${WANDB_NAME}
      entity: ${WANDB_ENTITY}
      save_dir: ./wandb
      log_model: false

[gabrieltseng]

Hi @robmarkcole - apologies for the delay.

Some thoughts:

• One possibility is that using the MSE loss is contributing to the blurred outputs. Have you tried (e.g.) a Huber loss?
• If possible, could you also share the validation R2 curves and MSE curves?
• I'd be curious about how a tiny (or base) model with ps 1 does - we are seeing that ps1 can occasionally lead to sharper predictions, so this may help too.

One another note (although I don't think this contributes to the blurring) - is it possible (especially with Base) to use a larger batch size than 1?

[robmarkcole]

• Typical curves:

• Try tiny with patch_size 1 -> predictions still blurry