Long runs with Convolutional CNP

[nmearl]

Thanks @wesselb for the email follow-up! I've taken your advice and made the following changes:

• Dynamic context/target split in training.

  photic/src/photic/run.py

  Line 78 in 4b55265

  # For every batch, flip a coin. If heads, do you what you now do. If tails,

• Random context size. Currently set to some random value between 0 the length of the light curve (200), or between 0 and a random split point.

  photic/src/photic/run.py

  Line 83 in 4b55265

  split = [np.random.randint(0, x.shape[2]), 35]

• Wrapping ar_predict with torch.no_grad did seem to help quite a bit. But I should also note that points_per_unit=1 seemed to also improve the memory usage as well. Seems to be working for now, but if the issues comes up again we can revisit.

I ended up running the convnp setup above for about 9h 33m 32s, and the results were not too bad. I show a gif below. The issue is mainly that after reaching peak, the predicted behavior becomes shallow and a little erratic. This isn't too much of a concern because we care less whether the prediction is accurate after the source has reached the peak, but is interesting.

I have also integrated with wandb -- very neat way to track the logs! I'm thinking now about what statistics to keep track of, and how we might want to quantify the predictive power of the NP, and will open another discussion to put my thoughts down.

I did have a quick question about the best way to save out the trained model. Currently, I'm just pickling the convcnp object, which seems to work fine, but if there's a preferred method, please let me know!

I've made quite a few changes to the repository so that it's more amenable to running on the compute clusters. I'll be setting off a run that I plan to let run for 12-24 hours. I do have a question, however, about how viable it is to resume training. That is, when submitting compute jobs, you can specify how long to use the resources, and, obviously, the more time you request, the longer you'll wait in queue. I'm curious if it's possible to request smaller blocks of time to continually train the same model, and how you might suggest doing so? Is it as simple as it seems -- i.e. just load the pickled model and continue training it as before?

[wesselb]

Hey @nmearl! Work craziness has finally passed, so I'm back again at it. :)

Your changes to the context/target split look good! I would perhaps even recommend to always use all remaining target points, since that should only be very minor increase in computational expense, but does give more signal:

        # For every batch, flip a coin. If heads, do you what you now do. If tails,
        # choose a point on the x-axis. Every point smaller than that point becomes
        # context, and every point bigger becomes target.
        if np.random.random() < 0.5:
            inds = np.random.permutation(x.shape[2])
            # Divide by two to ensure a sufficient number of target points.
            i_split = np.random.randint(0, x.shape[2] // 2)  
        else:
            inds = np.arange(x.shape[2])
            i_split = np.random.randint(0, x.shape[2])

        xc, yc = x[:, :, inds[: i_split]], y[:, :, inds[: i_split]]
        xt, yt = x[:, :, inds[i_split :]], y[:, :, inds[i_split :]]

Note that this only works if x is sorted, but, taking a quick look at your data generating functions, I think that should be right. I also added a // 2 to ensure that, in the first case of the if, there are always sufficient target points. What you ideally want is to vary the number of context points between zero and an upper bound, and keep the number of target points to a fixed, relatively high number. More target points is always better, though also more expensive.

I'm glad that ar_predict is now computationally more reasonable! :)

Looking at the W&B training loss, I'm a little concerned by that the loss jumps so much and appears to be so erratic. Have you checked whether the learning rate is right, and not possibly too high? What could be happening is that the loss values between the two cases of the if statement are of too different magnitude. For example, I would expect the loss to be lower for the np.random.permutation case than for the np.arange case. This difference in magnitude could interfere with the optimisation process. You might need some careful normalisation.

Looking at your sample predictions, this looks already quite good! I do think you should be able to do better and really optimise these predictions to perfection. (Perfection for you specific synthetic data; how it behaves on real data clearly is up in the air!) The way to further optimise the predictions is to make sure that what the model sees during training is as close as possible as what the model must do in practice. This was also the motivation for the idea of flipping a coin.

Currently, I'm just pickling the convcnp object, which seems to work fine, but if there's a preferred method, please let me know!

If that works fine, I would go for it! An alternative would be to save just the weights. Then, at loading time, you would first create a model with random weights and then load the saved weights. This approach is slightly safer, because you don't rely on pickle, which can be finicky. However, when you load the weights, you need to be absolutely sure that the model is exactly the same as the model you saved the weights for, otherwise you'll run into mismatch errors.

I'm curious if it's possible to request smaller blocks of time to continually train the same model, and how you might suggest doing so? Is it as simple as it seems -- i.e. just load the pickled model and continue training it as before?

This should certainly be possible. :) Indeed just load the model weights and continue training it as before. Note that the optimiser also has a state which you want to save and load when you resume training.