Hello,
I’m trying to use dropout at test-time with a neural network trained on MNIST, where the idea is to measure input-specific uncertainty. I do this by inputting a single test-set image, and having T models (defined by drop-out T times) make predictions, then I calculate the variance across the T model predictive probabilities for that class.
The problem is that when I make my predictions, all the T models output identical predictive probabilities for each of the 10 classes. I’m pretty sure that drop-out isn’t kicking in at test time. Can you please help?
I’m using Google Colab, running PyTorch version 1.4.0, what is automatically imported in Colab when I call ‘import torch’.
This is my network:
hyperparams = {'l1_out': 512,
'l2_out': 512,
'l1_drop': 0.5,
'l2_drop': 0.5,
'batch_size': 32,
'epochs': 10}
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(28 * 28, hyperparams['l1_out'])
self.dropout1 = nn.Dropout(hyperparams['l1_drop'])
self.fc2 = nn.Linear(hyperparams['l1_out'], hyperparams['l2_out'])
self.dropout2 = nn.Dropout(hyperparams['l2_drop'])
self.fc3 = nn.Linear(hyperparams['l2_out'], 10)
def forward(self, x):
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return F.log_softmax(x) #exponentiate logits to retrieve predicted probabilities
Now, I’ve decided to take T = 3, and to input a specific image. I’m also using two functions from two different sources right here on this forum, apply_dropout and enable_dropout.
train_loader = torch.utils.data.DataLoader(
datasets.MNIST('mnist-data/', train=True, download=True,
transform=transforms.Compose([transforms.ToTensor(),])),
batch_size=hyperparams['batch_size'], shuffle=True)
test_loader = torch.utils.data.DataLoader(
datasets.MNIST('mnist-data/', train=False, transform=transforms.Compose([transforms.ToTensor(),])
),
batch_size=hyperparams['batch_size'], shuffle=True)
#----------------------------------------------TRAIN ----------------------------------------------
net_trained = Net()
# create an Adam optimizer
optimizer = optim.Adam(net_trained.parameters(), lr=0.001)
# create a loss function
criterion = nn.CrossEntropyLoss()
epochs = hyperparams['epochs']
# run the main training loop
for epoch in range(epochs):
for batch_idx, (data, target) in enumerate(train_loader):
data, target = Variable(data), Variable(target)
# resize data from (batch_size, 1, 28, 28) to (batch_size, 28*28)
data = data.view(-1, 28*28)
optimizer.zero_grad()
net_out = net_trained(data)
loss = criterion(net_out, target)
loss.backward()
optimizer.step()
if batch_idx % 500 == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.data))
f = 'model'
torch.save(net_trained.state_dict(), f)
#----------------------------CHOOSE RANDOM IMAGE----------------------------------------------
for test_images, test_labels in test_loader:
sample_image = test_images[0]
sample_label = test_labels[0]
T = 3 # how many times to apply drop-out at test time
def apply_dropout(m):
for each_module in m.children():
if each_module.__class__.__name__.startswith('Dropout'):
each_module.train()
return m
def enable_dropout(m):
for each_module in m.modules():
if each_module.__class__.__name__.startswith('Dropout'):
each_module.train()
def uncertainties(p):
aleatoric = np.mean(p*(1-p), axis=0)
epistemic = np.mean(p**2, axis=0) - np.mean(p, axis=0)**2
return aleatoric, epistemic
#----------------------------PREDICT-------------------------------------------
def predict(model, image, label, T=T):
enable_dropout(model)
# prepare label
label = label.item()
# prepare image
image = image.flatten()
#image
standard_output = model(image)
_, standard_prediction = standard_output.max(0)
model.train()
# each vector will consist of T elements- the class-specific predictive probability from each model
zero_p_hat= []
one_p_hat= []
two_p_hat= []
three_p_hat= []
four_p_hat= []
five_p_hat= []
six_p_hat= []
seven_p_hat= []
eight_p_hat= []
nine_p_hat= []
# predict stochastic dropout model T times
for t in range(T):
enable_dropout(model) # STILL NOT WORKING WITH DROPOUT AT TEST-TIME
output = model(Variable(image))
output = output.detach()
output_prob = np.exp(output) #convert to predictive probabilities
zero_p_hat.append(output_prob[0].item()) # P( c = 0 | image)
one_p_hat.append(output_prob[1].item()) # P( c = 1 | image)
two_p_hat.append(output_prob[2].item()) # P( c = 2 | image)
three_p_hat.append(output_prob[3].item()) # P( c = 3 | image)
four_p_hat.append(output_prob[4].item()) # P( c = 4 | image)
five_p_hat.append(output_prob[5].item()) # P( c = 5 | image)
six_p_hat.append(output_prob[6].item()) # P( c = 6 | image)
seven_p_hat.append(output_prob[7].item()) # P( c = 7 | image)
eight_p_hat.append(output_prob[8].item()) # P( c = 8 | image)
nine_p_hat.append(output_prob[9].item()) # P( c = 9 | image)
# mean prediction
zero_var = np.var(zero_p_hat)
one_var = np.var(one_p_hat)
two_var = np.var(two_p_hat)
three_var = np.var(three_p_hat)
four_var = np.var(four_p_hat)
five_var = np.var(five_p_hat)
six_var = np.var(six_p_hat)
seven_var = np.var(seven_p_hat)
eight_var = np.var(eight_p_hat)
nine_var = np.var(nine_p_hat)
# will be a list of 10 elements, the variance of each class's predictions across T models
class_specific_uncertainties = [zero_var, one_var, two_var, three_var, four_var, five_var, six_var, seven_var, eight_var, nine_var]
variance_of_class_specific_predictions = class_specific_uncertainties[label]
# estimate uncertainties (eq. 4 )
# eq.4 in https://openreview.net/pdf?id=Sk_P2Q9sG
# see https://github.com/ykwon0407/UQ_BNN/issues/1
p_hat_lists = [zero_p_hat, one_p_hat, two_p_hat, three_p_hat, four_p_hat, five_p_hat, six_p_hat, seven_p_hat, eight_p_hat, nine_p_hat]
epistemic, aleatoric = uncertainties(np.array(p_hat_lists[label]))
return standard_prediction.item(), variance_of_class_specific_predictions, np.squeeze(aleatoric), np.squeeze(epistemic), p_hat_lists[label]
prediction, var_uncertainty, aleatoric, epistemic, _ = predict(net_untrained, sample_image, sample_label)
print("\n The model predicts: {} \n The ground truth is {}.\n With drop-out at test-time {} times, variance of class-specific predictions across the models is {}. \n Finally, aleatoric and epistemic uncertainties are {} and {}.\n The models give predictive probabilities: {}".format(prediction, sample_label, T, var_uncertainty, aleatoric, epistemic, _))
The output is:
The model predicts: 0
The ground truth is 0.
With drop-out at test-time 3 times, variance of class-specific predictions across the models is 0.0.
Finally, aleatoric and epistemic uncertainties are 0.0 and 0.013912441817748089.
The models give predictive probabilities:
[0.9858884215354919, 0.9858884215354919, 0.9858884215354919]