How to integrate the Pyro module with HuggingFace Transformers?
Author: h56cho
Pull request: no

Hello,
I am trying to convert the HuggingFace Transformer into a Bayesian neural network by using the `Pyro` module.
I provided my code below. Everything works well except I am stuck at the line `svi_loss = svi.step(input_ids = input_ids, attention_mask = attention_mask, labels = label)`. At that line an error is generated, because after converting the HuggingFace Transformer into a Pyro model, the new model does not have any set parameter (since it is a Bayesian model...so the weights for a Pyro model are not fixed, meaning the weights are sampled from a statistical distribution). Is there any way that I can get around this issue? I have also posted the similar question on Pyro forum. Thank you,

CODE: 

```python 
import torch
from torch import distributions
from transformers import RobertaTokenizer, RobertaForMultipleChoice, AdamW, get_constant_schedule
import pyro
import pyro.infer
import pyro.optim
import pyro.distributions as dist
import pyro.nn.module as module
import pyro.infer.autoguide.guides as guides
from torch import nn
from pyro.optim import Adam
from pyro.infer import SVI
from pyro.infer import Trace_ELBO
from pyro.infer import Predictive

# get the pre-trained HuggingFace RobertaForMultipleChoice and resize the token embeddings 
# after adding the special token
model_RobertaForMultipleChoice = RobertaForMultipleChoice.from_pretrained('Roberta-base')
        
# convert the HuggingFace model into a pyro model
module.to_pyro_module_(model_RobertaForMultipleChoice)

for m in model_RobertaForMultipleChoice.modules():
      for name, value in list(m.named_parameters(recurse=False)):
                setattr(m, name, module.PyroSample(prior=dist.Normal(0, 1)
                                             .expand(value.shape)
                                             .to_event(value.dim())))
        
        
# define parameters for training      
guide_delta = guides.AutoDelta(model_RobertaForMultipleChoice)
optimizer_2 = Adam({"lr": 0.000000055}) 
scheduler_2 = pyro.optim.StepLR({'optimizer': optimizer_2, 'optim_args': {'lr': 0.000000055}})
svi_delta = SVI(model_RobertaForMultipleChoice, guide_delta, optimizer_2, loss=Trace_ELBO())

# training loop
for m in range(num_iter):

        # calculate the loss and take a gradient step for svi
        # ERRORS OCCUR HERE 
        svi_loss = svi.step(input_ids = input_ids, 
                          attention_mask = attention_mask, 
                          labels = label)

        # update the with the calculated loss 
        total_svi_loss = total_svi_loss + svi_loss
       
        if m % log_interval == 0 and m > 0:
            cur_svi_loss = total_svi_loss / log_interval
            print('| epoch {:3d} | {:5d}/{:5d} batches | lr {:02.9f} | 
                   loss {:5.4f} | ppl {:8.4f}'.format(
                    epoch, m, int(num_lines_train/4), scheduler.get_lr()[0], 
                    cur_svi_loss, math.exp(cur_svi_loss)))
                   
            total_svi_loss = 0 
```

Comments:
please let me know how this progresses I am also interested in doing this

Hello,
Would it be possible for your team to also work on GPT2LMHeadModel on this same issue?

Thank you,

Hey @h56cho - did you find a solution here? Would be great if you can post code :-)  

Likewise, would be very interested in what came of this.