Add non-trivial convolutional encoder/decoder

[jemrobinson]

• Further simplify the naive linear encoder/decoder
• Add a non-linear CNN with activation functions

Testing

The following validation results come from the same minimal training (Dec 2019) and validation (Jan01-Jan15 2020) data.

I'm a bit worried about the fact that 500 epochs and 20 epochs look essentially the same, but perhaps this is due to the very small amount of training data?

Thoughts @marianovitasari20 @louisavz ?
Should we be doing some sort of normalisation at data-loading/data-prep time?

naive-null-naive: 781 params

  | Name      | Type               | Params | Mode
---------------------------------------------------------
0 | encoder_0 | NaiveLinearEncoder | 700    | train
1 | encoder_1 | NaiveLinearEncoder | 40     | train
2 | processor | NullProcessor      | 0      | train
3 | decoder   | NaiveLinearDecoder | 41     | train
---------------------------------------------------------
781       Trainable params
0         Non-trainable params
781       Total params
0.003     Total estimated model params size (MB)
14        Modules in train mode
0         Modules in eval mode

20 epochs

500 epochs

cnn-null-cnn: 1.8M params

  | Name      | Type          | Params | Mode
----------------------------------------------------
0 | encoder_0 | CNNEncoder    | 1.3 M  | train
1 | encoder_1 | CNNEncoder    | 1.4 K  | train
2 | processor | NullProcessor | 0      | train
3 | decoder   | CNNDecoder    | 937    | train
----------------------------------------------------
1.3 M     Trainable params
0         Non-trainable params
1.3 M     Total params
5.286     Total estimated model params size (MB)
95        Modules in train mode
0         Modules in eval mode

20 epochs

500 epochs

naive-unet-naive: 11M params

  | Name      | Type               | Params | Mode
---------------------------------------------------------
0 | encoder_0 | NaiveLinearEncoder | 700    | train
1 | encoder_1 | NaiveLinearEncoder | 40     | train
2 | processor | UNetProcessor      | 11.0 M | train
3 | decoder   | NaiveLinearDecoder | 41     | train
---------------------------------------------------------
11.0 M    Trainable params
0         Non-trainable params
11.0 M    Total params
43.901    Total estimated model params size (MB)
102       Modules in train mode
0         Modules in eval mode
https://wandb.ai/turing-seaice/leaderboard/runs/aizqh771

20 epochs

500 epochs

cnn-unet-cnn: 12.8M params

  | Name      | Type          | Params | Mode
----------------------------------------------------
0 | encoder_0 | CNNEncoder    | 1.3 M  | train
1 | encoder_1 | CNNEncoder    | 1.4 K  | train
2 | processor | UNetProcessor | 11.0 M | train
3 | decoder   | CNNDecoder    | 937    | train
----------------------------------------------------
12.3 M    Trainable params
0         Non-trainable params
12.3 M    Total params
49.184    Total estimated model params size (MB)
183       Modules in train mode
0         Modules in eval mode

20 epochs

500 epochs

[github-actions]

Coverage report

Click to see where and how coverage changed

File	Statements	Missing	Coverage	Coverage (new stmts)	Lines missing
ice_station_zebra/models/common
__init__.py
conv_block_downsample.py
conv_block_upsample.py
resizing_average_pool_2d.py
ice_station_zebra/models/decoders
__init__.py
base_decoder.py
cnn_decoder.py
naive_linear_decoder.py
ice_station_zebra/models/encoders
__init__.py
base_encoder.py
cnn_encoder.py
naive_linear_encoder.py
Project Total

This report was generated by python-coverage-comment-action

[jemrobinson]

Latent space visualisation

I've done some quick visualisations of the latent space for NaiveLinear and CNN encoders.

Naive linear

CNN

Normalisation

Now we can see that the CNN has actually lost some information during the encoding process, let's try adding a normalisation layer as the first step of the encoder.

CNN with initial BatchNorm

CNN with initial LayerNorm

CNN with initial InstanceNorm2d

[jemrobinson]

After adding a normalisation layer and running for 20 epochs

cnn-unet-cnn

naive-unet-naive

[IFenton]

Looking at this, I'm wondering about whether we need to do any regridding / reprojecting. At the moment, the SIC is centred round the south pole (-90 lat, 0 long), but I assume the ERA5 data is centred on something like -45 lat, 0 long?

[IFenton]

Reviewing this

[jemrobinson]

@IFenton: I could see that the different scales/grids for ERA5 and OSISAF might mean that we end up depending more heavily on OSISAF but not that we can't produce anything at all. I wonder whether there's an additional issue in the decoder?

[IFenton]

@jemrobinson Curiously, I've just done a couple of runs of naive_unet_naive, and it's producing pretty good results, e.g. https://wandb.ai/turing-seaice/leaderboard/runs/3g3iduk3?nw=nwuserifenton. It's just running on the defaults, so not quite sure what's going on

[IFenton]

LGTM

[IFenton]

Suggested change

	name: encode-ddpm-decode
	name: cnn-ddpm-cnn