How to fine-tune the output head of the Pyro Neural Network?

Misc.
1

Hello,
I have a Pyro neural network model that only has Bayesian layer for its output head. The main body of this neural network has pre-trained frequentist weights, and I am not intending to change them much. I just want to train the guide that I use for the output head of my neural network, while fine-tuning the main body frequentist weights as necessary.

How can I achieve this with Pyro models?

right now I am doing:

optimizer_args = {'lr': 0.000010}
optimizer_3 = torch.optim.Adam
scheduler_args = {'optimizer': optimizer_3,
                          'step_size' : 1, 'gamma' : 1.5, 
                          'optim_args' : optimizer_args}
scheduler_3 = pyro.optim.StepLR(scheduler_args)

svi_diag_normal = SVI(model, guide_diag_normal, scheduler_3,
                              loss=TraceEnum_ELBO(max_plate_nesting=0))

# train the guide of the output head
train(model, input_data, optimizer, scheduler, log_interval,
                      svi, guide, epoch)

But this applies the learning rate of 0.000010 to both frequentist weights and the guide for the output head. (am I correct? Please correct me if this is not true).

How should I fine-tune my pyro neural network?

Thank you,

2

there are different ways to do this but one way is to use torch optimizers and take control of the process yourself: Customizing SVI objectives and training loops — Pyro Tutorials 1.8.4 documentation