openclimatefix / satflow

Satellite Optical Flow with machine learning models
https://satflow.readthedocs.io/en/stable/
MIT License
61 stars 10 forks source link

Add Pytorch Lightning Training #15

Closed jacobbieker closed 3 years ago

jacobbieker commented 3 years ago

Relates to #14, finishing up things from #13 for actually having a working training loop and first model.

jacobbieker commented 3 years ago

The recurrent UNET seemed to have some issues, now just trying with a simply ConvLSTM encoder-decoder, which is training and seems to work fine

jacobbieker commented 3 years ago

Trains at 7.5sec/iteration (batch size of 8) with 8 dataloader workers off a 7200rpm hard drive on RTX 3070 laptop GPU, so definitely not the fastest, but does okay for now.