def test_half_cauchy_prior_log_prob_log_transform(self, cuda=False):
device = torch.device("cuda") if cuda else torch.device("cpu")
prior = HalfCauchyPrior(0.1, transform=torch.exp)
dist = HalfCauchy(0.1)
t = torch.tensor(0.0, device=device)
self.assertTrue(torch.equal(prior.log_prob(t), dist.log_prob(t.exp())))
t = torch.tensor([-1, 0.5], device=device)
self.assertTrue(torch.equal(prior.log_prob(t), dist.log_prob(t.exp())))
t = torch.tensor([[-1, 0.5], [0.1, -2.0]], device=device)
self.assertTrue(torch.equal(prior.log_prob(t), dist.log_prob(t.exp())))
def test_half_cauchy_prior_log_prob(self, cuda=False):
device = torch.device("cuda") if cuda else torch.device("cpu")
prior = HalfCauchyPrior(0.1)
dist = HalfCauchy(0.1)
t = torch.tensor(1.0, device=device)
self.assertTrue(torch.equal(prior.log_prob(t), dist.log_prob(t)))
t = torch.tensor([1.5, 0.5], device=device)
self.assertTrue(torch.equal(prior.log_prob(t), dist.log_prob(t)))
t = torch.tensor([[1.0, 0.5], [3.0, 0.25]], device=device)
self.assertTrue(torch.equal(prior.log_prob(t), dist.log_prob(t)))
def test_half_cauchy_prior_batch_log_prob(self, cuda=False):
device = torch.device("cuda") if cuda else torch.device("cpu")
prior = HalfCauchyPrior(0.1)
dist = HalfCauchy(0.1)
t = torch.ones(2, device=device)
self.assertTrue(torch.equal(prior.log_prob(t), dist.log_prob(t)))
t = torch.ones(2, 2, device=device)
self.assertTrue(torch.equal(prior.log_prob(t), dist.log_prob(t)))
scale = torch.tensor([0.1, 1.0], device=device)
prior = HalfCauchyPrior(scale)
dist = HalfCauchy(scale)
t = torch.ones(2, device=device)
self.assertTrue(torch.equal(prior.log_prob(t), dist.log_prob(t)))
t = torch.ones(2, 2, device=device)
self.assertTrue(torch.equal(prior.log_prob(t), dist.log_prob(t)))
with self.assertRaises(ValueError):
prior.log_prob(torch.ones(3, device=device))
with self.assertRaises(ValueError):
prior.log_prob(torch.ones(2, 3, device=device))
def expand(self, batch_shape):
return HalfCauchy(self.loc.expand(batch_shape), self.scale.expand(batch_shape))
def __init__(self, scale, validate_args=None, transform=None):
TModule.__init__(self)
HalfCauchy.__init__(self, scale=scale, validate_args=validate_args)
self._transform = transform
def rsample(self, sample_shape=torch.Size([])):
local_shrinkage = HalfCauchy(1).rsample(self.scale.shape)
param_sample = Normal(0, local_shrinkage * self.scale).rsample(sample_shape)
return param_sample
def test_half_cauchy_prior_to_gpu(self):
if torch.cuda.is_available():
prior = HalfCauchy(1.0).cuda()
self.assertEqual(prior.concentration.device.type, "cuda")
self.assertEqual(prior.rate.device.type, "cuda")
def __init__(self, in_features, out_features, parameters, device):
"""
Args:
in_features: int, number of input features
out_features: int, number of output features
parameters: instance of class HorseshoeHyperparameters
device: cuda device instance
"""
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.device = device
# Scale to initialize weights, according to Yingzhen's work
if parameters.horseshoe_scale == None:
scale = 1. * np.sqrt(6. / (in_features + out_features))
else:
scale = parameters.horseshoe_scale
# Initialization of parameters of prior distribution
# weight parameters
self.prior_tau_shape = torch.Tensor([0.5])
# local shrinkage parameters
self.prior_lambda_shape = torch.Tensor([0.5])
self.prior_lambda_rate = torch.Tensor(
[1 / parameters.weight_cauchy_scale**2])
# global shrinkage parameters
self.prior_v_shape = torch.Tensor([0.5])
self.prior_theta_shape = torch.Tensor([0.5])
self.prior_theta_rate = torch.Tensor(
[1 / parameters.global_cauchy_scale**2])
# Initialization of parameters of variational distribution
# weight parameters
self.beta_mean = nn.Parameter(
torch.Tensor(out_features, in_features).uniform_(-scale, scale))
self.beta_rho = nn.Parameter(
torch.ones([out_features, in_features]) *
parameters.beta_rho_scale)
self.beta = ReparametrizedGaussian(self.beta_mean, self.beta_rho)
# local shrinkage parameters
self.lambda_shape = self.prior_lambda_shape * torch.ones(in_features)
self.lambda_rate = self.prior_lambda_rate * torch.ones(in_features)
self.lambda_ = InverseGamma(self.lambda_shape, self.lambda_rate)
# Sample from half-Cauchy to initialize the mean of log_tau
# We initialize the parameters using a half-Cauchy because this
# is the prior distribution over tau
if parameters.log_tau_mean == None:
distr = HalfCauchy(1 / np.sqrt(self.prior_lambda_rate))
sample = distr.sample(torch.Size([in_features])).squeeze()
self.log_tau_mean = nn.Parameter(torch.log(sample))
else:
self.log_tau_mean = parameters.log_tau_mean
self.log_tau_rho = nn.Parameter(
torch.ones(in_features) * parameters.log_tau_rho_scale)
self.log_tau = ReparametrizedGaussian(self.log_tau_mean,
self.log_tau_rho)
# bias parameters
self.bias_mean = nn.Parameter(torch.zeros([1, out_features], ))
self.bias_rho = nn.Parameter(
torch.ones([1, out_features]) * parameters.bias_rho_scale)
self.bias = ReparametrizedGaussian(self.bias_mean, self.bias_rho)
# global shrinkage parameters
self.theta_shape = self.prior_theta_shape
self.theta_rate = self.prior_theta_rate
self.theta = InverseGamma(self.theta_shape, self.theta_rate)
# Sample from half-Cauchy to initialize the mean of log_v
# We initialize the parameters using a half-Cauchy because this
# is the prior distribution ovev
if parameters.log_v_mean == None:
distr = HalfCauchy(1 / np.sqrt(self.prior_theta_rate))
sample = distr.sample()
self.log_v_mean = nn.Parameter(torch.log(sample))
else:
self.log_v_mean = parameters.log_v_mean
self.log_v_rho = nn.Parameter(
torch.ones([1, 1]) * parameters.log_v_rho_scale)
self.log_v = ReparametrizedGaussian(self.log_v_mean, self.log_v_rho)