def _setup_optimizer(self, optimizer, lr, lr_decay):
if optimizer == 'Adam':
self.optimizer = torch.optim.Adam(self.posterior.parameters(), lr=lr)
elif optimizer == 'SGD':
self.optimizer = torch.optim.SGD(self.posterior.parameters(), lr=lr)
else:
raise NotImplementedError('Optimizer must be Adam or SGD')
if lr_decay < 1.0:
self.lr_scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, 1000, gamma=lr_decay)
else:
self.lr_scheduler = DummyLRScheduler()
def _setup_optimizer(self, optimizer, lr, lr_decay):
assert hasattr(
self, 'particles'
), "SVGD must be initialized before setting up optimizer"
if optimizer == 'Adam':
self.optimizer = torch.optim.Adam([self.particles], lr=lr)
elif optimizer == 'SGD':
self.optimizer = torch.optim.SGD([self.particles], lr=lr)
else:
raise NotImplementedError('Optimizer must be Adam or SGD')
if lr_decay < 1.0:
self.lr_scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer,
1000,
gamma=lr_decay)
else:
self.lr_scheduler = DummyLRScheduler()