def test_freeze_module(self):
#create module, freeze
#check if two inferences keep equal
blinear = BayesianEmbedding(10, 10)
to_feed = torch.Tensor([i for i in range(9)]).unsqueeze(0).long()
self.assertEqual((blinear(to_feed) != blinear(to_feed)).any(), torch.tensor(True))
frozen_feedforward = blinear.forward_frozen(to_feed)
blinear.freeze = True
self.assertEqual((blinear.forward(to_feed) == frozen_feedforward).all(), torch.tensor(True))
def test_sequential_cpu(self):
#check if we can create sequential models chaning our Bayesian Linear layers
model = nn.Sequential(BayesianEmbedding(10, 10),
nn.LSTM(10, 15),)
to_feed = torch.Tensor([i for i in range(9)]).unsqueeze(0).long()
#if this works, the test will pass
result = model(to_feed)
pass
def test_variational_inference(self):
#create module, check if inference is variating
blinear = BayesianEmbedding(10, 10)
linear = nn.Embedding(10, 10)
to_feed = torch.Tensor([i for i in range(9)]).unsqueeze(0).long()
self.assertEqual((blinear(to_feed) != blinear(to_feed)).any(), torch.tensor(True))
self.assertEqual((linear(to_feed) == linear(to_feed)).all(), torch.tensor(True))
pass
def test_sequential_cuda(self):
#check if we can create sequential models chaning our Bayesian Linear layers
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = nn.Sequential(BayesianEmbedding(10, 10),
nn.LSTM(10, 15),).to(device)
to_feed = torch.Tensor([i for i in range(9)]).unsqueeze(0).long().to(device)
#if this works, the test will pass
result = model(to_feed)
pass
def test_kl_divergence(self):
#create model, sample weights
#check if kl divergence between apriori and a posteriori is working
blinear = BayesianEmbedding(10, 10)
to_feed = torch.Tensor([i for i in range(9)]).unsqueeze(0).long()
predicted = blinear(to_feed)
complexity_cost = blinear.log_variational_posterior - blinear.log_prior
self.assertEqual((complexity_cost == complexity_cost).all(), torch.tensor(True))
pass
def test_infer_shape_1_sample(self):
blinear = BayesianEmbedding(10, 10)
linear = nn.Embedding(10, 10)
to_feed = torch.Tensor([i for i in range(9)]).unsqueeze(0).long()
b_infer = linear(to_feed)
l_infer = blinear(to_feed)
self.assertEqual(b_infer.shape, l_infer.shape)
self.assertEqual((b_infer == b_infer).all(), torch.tensor(True))
pass
def convert_layer_to_bayesian(layer):
if isinstance(layer, torch.nn.Linear):
new_layer = BayesianLinear(layer.in_features, layer.out_features)
elif isinstance(layer, nn.Embedding):
new_layer = BayesianEmbedding(layer.num_embeddings,
layer.embedding_dim)
elif isinstance(layer, nn.Conv1d):
new_layer = BayesianConv1d(
layer.in_channels,
layer.out_channels,
kernel_size=layer.kernel_size[0],
groups=layer.groups,
padding=layer.padding,
dilation=layer.dilation,
)
else:
Warning(
f"Could not find correct type for conversion of layer {layer} with type {type(layer)}"
)
new_layer = layer
return new_layer
def test_init_bayesian_layer(self):
module = BayesianEmbedding(10, 10)
pass
def test_inheritance(self):
#check if bayesian linear has nn.Module and BayesianModule classes
blinear = BayesianEmbedding(10, 10)
self.assertEqual(isinstance(blinear, (nn.Module)), True)
self.assertEqual(isinstance(blinear, (BayesianModule)), True)