SGHMC - does my code do what I think it does?

username: PhysicsIsFun
cooked: <p>Greetings,</p>
<p>I try to implement Stochastic Gradient Hamiltonian Monte Carlo (SGHMC) scheme to sample from the Bayesian posterior of neural networks.</p>
<p>I think I implemented everything correctly, but I am a bit worried that PyTorch does stuff “under the hood” that I did not intend.</p>
<p>The update rule of the scheme is given by<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/c/ac05cfe985254d7e4c32b2db885dcefcdafdf733.png" data-download-href="https://discuss.pytorch.org/uploads/default/ac05cfe985254d7e4c32b2db885dcefcdafdf733" title="pytorch1"><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/c/ac05cfe985254d7e4c32b2db885dcefcdafdf733.png" alt="pytorch1" data-base62-sha1="oxMzNy3cMVdF1lL5IJDqrilJ2nN" width="690" height="106" data-dominant-color="F5F5F5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">pytorch1</span><span class="informations">1008×156 7.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
where the R_i is a standard normal random variable.<br>
The \theta_i are the network parameters, and the p_i corresponding momenta (in the code, I store them in a <code>buf</code> variable within the parameters.<br>
The nabla U is the stochastic gradient of the loss plus weight decay.</p>
<p>My code assumes a given network <code>model</code> with an <code>evaluate</code> function that computes losses and accuracies of a given data loader.</p>
<p>Does this class look correct?</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import numpy as np
import time
import copy
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  ## use GPU if available

class SGHMC(nn.Module):
    
    def __init__(self, model, train_loader, test_loader, criterion, lr, weight_decay, alpha, epochs):
        super(SGHMC, self).__init__()
        self.model = model
        self.train_loader = train_loader
        self.test_loader = test_loader
        self.criterion = criterion
        self.lr = lr
        self.weight_decay = weight_decay
        self.alpha = alpha
        self.epochs = epochs
        
        
    def train(self):
        
        loss_train = np.zeros(self.epochs+1)                 
        accu_train = np.zeros(self.epochs+1)
        loss_test = np.zeros(self.epochs+1)
        accu_test = np.zeros(self.epochs+1)
        
        (loss_train[0], accu_train[0]) = self.model.evaluate(self.train_loader)
        (loss_test[0], accu_test[0]) = self.model.evaluate(self.test_loader)        
        
        datasize = len(self.train_loader.dataset)
        
        squeeze = True if type(self.criterion) == torch.nn.modules.loss.BCELoss else False   # squeeze network output
                                                                                             # for BCELoss (not required
                                                                                             # for NLLLoss)
        # initialize momenta
        for p in self.model.parameters():
            p.buf = torch.randn(p.size()).to(device)

        # train routine
        for epoch in range(1, self.epochs+1):
            self.model.train()
            
            for batch_idx, (data, target) in enumerate(self.train_loader):
                data, target = data.to(device), target.to(device)
                self.model.zero_grad()
                output = self.model(data) 
                if squeeze: output=output.squeeze()
                loss = self.criterion(output, target)*datasize
                loss.backward()
                self.update_params()
        
            (loss_train[epoch], accu_train[epoch]) = self.model.evaluate(self.train_loader)
            (loss_test[epoch], accu_test[epoch]) = self.model.evaluate(self.test_loader)
          
        return (loss_train, loss_test, accu_train, accu_test)


    def update_params(self):

        for p in self.model.parameters():
        
            p.grad.data.add_(p.data, alpha=self.weight_decay)  # adding weight decay to gradients
            
            # update momenta
            eps = torch.randn(p.size()).to(device)  
            p.buf.mul_(1-self.alpha)
            p.buf.add_(-self.lr*p.grad.data + (2.0 * self.lr * self.alpha)**.5 * eps)
            
            # update parameter
            p.data.add_(p.buf)
</code></pre>