Two simple implementations of sg-mcmc

sg-mcmc.py

from torch.optim.optimizer import Optimizer
class OptimizerSGLD(Optimizer):
    def __init__(self, net, alpha=1e-4, sgmcmc=True):
        super(OptimizerSGLD, self).__init__(net.parameters(), {})
        self.net = net
        self.sgmcmc = sgmcmc
        self.alpha = alpha
        self.noise_std = (2*self.alpha)**0.5
    
    @torch.no_grad()
    def step(self):
        for name, par in self.net.named_parameters():
            newpar = par - self.alpha*par.grad
            if self.sgmcmc:
                noise = torch.normal(torch.zeros_like(par), std=self.noise_std)
                newpar += noise
            par.copy_(newpar)

class OptimizerSGHMC(Optimizer):
    def __init__(self, net, alpha=1e-4, nu=1.0, sgmcmc=True):
        super(OptimizerSGHMC, self).__init__(net.parameters(), {})
        self.net = net
        self.sgmcmc = sgmcmc
        self.alpha = alpha
        self.nu = nu
        self.noise_std = (2*self.alpha*self.nu)**0.5

        self.momentum = {key : torch.zeros_like(par) for key, par in self.net.named_parameters()}
    
    @torch.no_grad()
    def step(self):
        for name, par in self.net.named_parameters():
            newpar = par + self.momentum[name]
            par.copy_(newpar)

            # Update momentum par:
            self.momentum[name] = (1-self.nu)*self.momentum[name] - self.alpha*par.grad
            if self.sgmcmc:
                noise = torch.normal(torch.zeros_like(self.momentum[name]), std=self.noise_std)
                self.momentum[name] += noise