search

openclimatefix / graph_weather

PyTorch implementation of Ryan Keisler's 2022 "Forecasting Global Weather with Graph Neural Networks" paper (https://arxiv.org/abs/2202.07575)
MIT License
188 stars 47 forks source link

Running pl_graph_weather.py in graph_weather/train causes "Killed" #35

Closed paapu88 closed 1 year ago

paapu88 commented 2 years ago

Dear Developers, Wanted to test how training goes 'out of the box' with graph_weather. First problem was that I ran out of disk space, but thats ok now.

Other problem (not fixed) is that I get 'Killed' and program stops (see below). Probably one has run out of memory. Is there any way to fix this? Typically one sets smaller batch size, but I think it is already smallelst possible (1?).

My machine is
nvidia tesla v100 product: GV100GL [Tesla V100 SXM2 16GB]

Terveisin, Markus

ubuntu@gpumok:~/git/graph_weather/train$ python3 pl_graph_weather.py 
/usr/lib/python3/dist-packages/pkg_resources/__init__.py:116: PkgResourcesDeprecationWarning: 0.1.43ubuntu1 is an invalid version and will not be supported in a future release
  warnings.warn(
/usr/lib/python3/dist-packages/pkg_resources/__init__.py:116: PkgResourcesDeprecationWarning: 1.1build1 is an invalid version and will not be supported in a future release
  warnings.warn(
Using custom data configuration openclimatefix--gfs-surface-pressure-2deg-ca6875a1d925d744
Reusing dataset parquet (/data10t/huggingface/datasets/openclimatefix___parquet/openclimatefix--gfs-surface-pressure-2deg-ca6875a1d925d744/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec)
Using custom data configuration openclimatefix--gfs-surface-pressure-2deg-ca6875a1d925d744
Reusing dataset parquet (/data10t/huggingface/datasets/openclimatefix___parquet/openclimatefix--gfs-surface-pressure-2deg-ca6875a1d925d744/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec)
#1: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:14:22<00:00, 282.63s/ex]
#7: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:14:23<00:00, 282.65s/ex]
#2: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:14:29<00:00, 282.77s/ex]
#9: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:14:35<00:00, 282.89s/ex]
#0: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:14:42<00:00, 283.02s/ex]
#4: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:14:50<00:00, 283.15s/ex]
#14: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:14:59<00:00, 283.33s/ex]
#11: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:01<00:00, 283.36s/ex]
#8: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:08<00:00, 283.49s/ex]
#3: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:09<00:00, 283.51s/ex]
#10: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:10<00:00, 283.54s/ex]
#5: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:15<00:00, 283.61s/ex]
#12: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:23<00:00, 283.76s/ex]
#6: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:24<00:00, 283.78s/ex]
#13: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:31<00:00, 283.92s/ex]
#15: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:31<00:00, 283.92s/ex]
  0%|                                                                                                                                                            | 0/1 [00:00<?, ?ba/s]Killed
Killed
ubuntu@gpumok:~/git/graph_weather/train$ ███████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:01<00:00, 279.11s/ex]
#6: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:24<00:00, 272.73s/ex]
#10: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:10<00:00, 279.13s/ex]
#12: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:23<00:00, 272.14s/ex]
#13: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:31<00:00, 264.74s/ex]
#15: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:31<00:00, 270.98s/ex]
jacobbieker commented 2 years ago

Hi, yeah, it uses a lot of ram by default, as its using around 630 variables, vs the 72 in the paper. The easiest way would be to, in the get_mean_std bit, comment out some of the variables you are not interested in, then it should only use that subset of all of them. The training script works on my laptop, with 64GB ram and a 8GB GPU, but just barely. A more involved, but potentially less memory intensive option is to load the original Zarr files directly from https://huggingface.co/datasets/openclimatefix/gfs-reforecast and only select the variables you are interested in. That should result in a lot smaller memory footprint. But if you want to use all the variables in the NWPs, then it just uses a lot of memory.