I want to use a Bayesian Neural Network for a regression task. To do that I converted a BNN from this paper to Python 3. The provided training script runs and I receive a pickle file, which I want to use to predict a value for my regression. Event though the loss for the training doesn't really go down, but this seems also be the case for the models used in the paper. The input of my network is a vector with 59 features, where one feature is my target variable. It's pretty similar to the boston housing dataset. The MLP inside the BNN returns two vales: $\mu$ and $\sigma$. If I understand BNNs correctly $\mu$ and $\sigma$ refer to the mean and standard deviation of my input and I should use $\mu$ to predict my values.
But when I use input vector I always receive a value that is roughly around the mean of my target variable. So lets say the mean is 62. When I insert a vector from my test set, that should return the target variable 142, I will get 61. When I do the same with a vector, that should return 21, I will get 63. So my predicted values are always around the mean of the target variable.
What do I have to do, to get real predictions from my input data? I never worked with BNNs before and it's quite hard to find any resources on how I should model my network. Maybe someone has a nice tutorial to learn how I can use this model to predict my target variable. I have to use this class and can't use Pyro or any other framework currently, because otherwise I can't use the other components from the paper.
To predict the values I used the predict() method from the converted BNN Class:
class GaussianBNN(BaseNet):
def __init__(self, model, N_train, lr=1e-2, cuda=True, eps=1e-3, grad_std_mul=20):
super(GaussianBNN, self).__init__()
self.lr = lr
self.model = model
self.cuda = cuda
self.N_train = N_train
self._create_network()
self._create_optimizer()
self.schedule = None # [] #[50,200,400,600]
self.epoch = 0
self.grad_buff = []
self.grad_std_mul = grad_std_mul
self.max_grad = 1e20
self.eps = eps
self.weight_set_samples = []
def _create_network(self):
torch.manual_seed(42)
if self.cuda:
torch.cuda.manual_seed(42)
if self.cuda:
self.model.cuda()
cudnn.benchmark = True
print('Total params: %.2fM' % (self.get_nb_parameters() / 1000000.0))
def _create_optimizer(self):
self.optimizer = StochasticHamiltonMonteCarloSampler(params=self.model.parameters(), lr=self.lr, base_C=0.05, gauss_sig=0.1)
def fit(self, x, y, burn_in=False, resample_momentum=False, resample_prior=False):
self.set_mode_train(train=True)
x, y = variable_to_tensor_list(variables=(x, y), cuda=self.cuda)
self.optimizer.zero_grad()
mu, sigma = self.model(x)
sigma = sigma.clamp(min=self.eps)
loss = -diagonal_gauss_loglike(y, mu, sigma).mean(dim=0) * self.N_train
loss.backward()
if len(self.grad_buff) > 100:
self.max_grad = np.mean(self.grad_buff) + self.grad_std_mul * np.std(self.grad_buff)
self.grad_buff.pop(0)
self.grad_buff.append(nn.utils.clip_grad_norm_(parameters=self.model.parameters(),
max_norm=self.max_grad, norm_type=2))
if self.grad_buff[-1] >= self.max_grad:
print(self.max_grad, self.grad_buff[-1])
self.grad_buff.pop()
self.optimizer.step(burn_in=burn_in, resample_momentum=resample_momentum, resample_prior=resample_prior)
return loss * x.shape[0] / self.N_train, mu, sigma
def eval(self, x, y):
self.set_mode_train(train=False)
x, y = variable_to_tensor_list(variables=(x, y), cuda=self.cuda)
mu, sigma = self.model(x)
sigma = sigma.clamp(min=self.eps)
loss = -diagonal_gauss_loglike(y, mu, sigma).mean(dim=0) * self.N_train
return loss * x.shape[0] / self.N_train, mu, sigma
@staticmethod
def unnormalised_eval(pred_mu, pred_std, y, y_mu, y_std, gmm=False):
ll = gaussian_mixture_model_loglike(pred_mu, pred_std, y, y_mu, y_std, gmm=gmm) # this already computes sum
if gmm:
pred_mu = pred_mu.mean(dim=0)
rms = get_root_mean_square(pred_mu, y, y_mu, y_std) # this already computes sum
return rms, ll
def predict(self, x):
self.set_mode_train(train=False)
x, = variable_to_tensor_list(variables=(x,), cuda=self.cuda)
mu, sigma = self.model(x)
return mu, sigma
def save_sampled_net(self, max_samples):
if len(self.weight_set_samples) >= max_samples:
self.weight_set_samples.pop(0)
self.weight_set_samples.append(copy.deepcopy(self.model.state_dict()))
logger.warning(f"Saving Samples: {len(self.weight_set_samples)} for Max Samples: {max_samples}")
logger.warning(f"Samples: {self.weight_set_samples}")
return None
def sample_predict(self, x, num_samples, grad=False):
self.set_mode_train(train=False)
if num_samples == 0:
num_samples = len(self.weight_set_samples)
x, = variable_to_tensor_list(variables=(x,), cuda=self.cuda)
if grad:
self.optimizer.zero_grad()
if not x.requires_grad:
x.requires_grad = True
mu_vec = x.new(num_samples, x.shape[0], self.model.output_dim)
std_vec = x.new(num_samples, x.shape[0], self.model.output_dim)
# iterate over all saved weight configuration samples
# TODO: mu_vec and std_vec are changed here, prove that the value isnt changed inplace
for idx, weight_dict in enumerate(self.weight_set_samples):
if idx == num_samples:
break
self.model.load_state_dict(weight_dict)
mu, std = self.model(x)
mu_vec[idx] = mu.detach().clone()
std_vec[idx] = std.detach().clone()
if grad:
# TODO: check if same things are returner here
return mu_vec[:idx], std_vec[:idx]
else:
return mu_vec[:idx], std_vec[:idx]
def get_weight_samples(self, num_samples=0):
weight_vec = []
if num_samples == 0 or num_samples > len(self.weight_set_samples):
num_samples = len(self.weight_set_samples)
for idx, state_dict in enumerate(self.weight_set_samples):
if idx == num_samples:
break
for key in state_dict.keys():
if 'weight' in key:
weight_mtx = state_dict[key].cpu()
for weight in weight_mtx.view(-1):
weight_vec.append(weight)
return np.array(weight_vec)
def save_weights(self, filename):
save_object(self.weight_set_samples, filename)
def load_weights(self, filename):
self.weight_set_samples = load_object(filename)
The referenced model here is the following MLP:
class GaussianMLP(nn.Module):
def __init__(self, input_dim, width, depth, output_dim, flatten_image):
super(GaussianMLP, self).__init__()
self.input_dim = input_dim
self.output_dim = output_dim
self.width = width
self.depth = depth
self.flatten_image = flatten_image
layers = [nn.Linear(input_dim, width), nn.ReLU()]
for i in range(depth - 1):
layers.append(nn.Linear(width, width))
layers.append(nn.ReLU())
layers.append(nn.Linear(width, 2 * output_dim))
self.block = nn.Sequential(*layers)
def forward(self, x):
if self.flatten_image:
x = x.view(-1, self.input_dim)
x = self.block(x)
mu = x[:, :self.output_dim]
sigma = F.softplus(x[:, self.output_dim:])
return mu, sigma
You can find the training script here: [LINK]