SimonTopp
commented
2 years ago Wanted to update this quickly. At this point I don't this this is explicitly an issue with the GW loss function. I was able to look at when it breaks, and it seems like it does so when the predicted temperatures are all Nan. Why sometimes the predicted temperatures are all Nan and why the plain mutlitask rmse can handle this is another question. These Nan predictions are seemingly random, and we've seen them everywhere from the third epoch to the 100th (although they do seem to occur more frequently later in training). I'm digging into it more, but I basically have to let it run and hope it breaks before I can step through all the code and trace things back, so it's a slow process.
aappling-usgs
janetrbarclay
The groundwater loss function is currently explicitly put onto a CPU regardless of GPU availability because for some reason running it on a GPU causes the loss to go to Nan within the ~1st epoch. Not totally sure why this is, but based on various threads (i.e. keras-team/keras#1244), it could be either a hardware or software issue.
Things we've tried, but that have been unsuccessful, include gradient clipping (norm and value), increasing epsilon in the optimizer, changing the optimizer, limiting memory growth by the GPU, swapping out all Nans in Delta Phi and AR to 0, and changing all our values from float64 to float32. The loss function handles pre-training just fine, which might suggest it has something to do with too many Nans in our fine tuning dataset? For now the loss function is stable regardless of where you run it, but solving this eventually would be good.