[Python] .grad attribute of a Tensor that is not a leaf Tensor is being accessed

I am trying to implement a dictionary learning algorithm and have been struggling with the following error.

UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations. (Triggered internally at aten/src/ATen/core/TensorBody.h:417.)

I know this is a warning, but since I need the gradient later, not calculating the gradient ends up throwing a NoneType error at the following line in my code: P2 = -0.5 * (gradient / torch.norm(gradient, dim=0)) + P1 This is in a method to calculate the step to take:

def get_spherical_step(self, start, gradient, step_size):
with torch.no_grad():
P1 = start / torch.norm(start, dim=0)
P2 = -0.5 * (gradient / torch.norm(gradient, dim=0)) + P1
P2 /= torch.norm(P2, dim=0)

projection_p1_p2 = (P1 * P2).sum(dim=0, keepdim=True) * P1
orthogonal_part = P2 - projection_p1_p2

end = P1 * math.cos(step_size) + (orthogonal_part / torch.norm(orthogonal_part, dim=0, keepdim=True)) * math.sin(step_size)

epsilon = 1e-7
zero_gradient_mask = (torch.norm(gradient, dim=0) <= epsilon) | (torch.norm(orthogonal_part, dim=0) <= epsilon)
end[:, zero_gradient_mask] = P1[:, zero_gradient_mask]

return end


This is the method that takes that step:

def optimizer_step(
self
,
batch
,
loss_function
):

if
self.current_probe_step == self.max_probe_steps:
self.reset_probe()

self.current_probe_step += 1


with
torch.no_grad():
smaller_step_R = torch.linalg.lstsq(self.smaller_step_dictionary, batch).solution
normal_step_R = torch.linalg.lstsq(self.dictionary, batch).solution
bigger_step_R = torch.linalg.lstsq(self.bigger_step_dictionary, batch).solution

dictionaries = [self.smaller_step_dictionary, self.dictionary, self.bigger_step_dictionary]
step_sizes = [self.step_size / 2, self.step_size, self.step_size * 2]

batch_losses = []

for
i, dictionary
in
enumerate(dictionaries):
dictionary.requires_grad_(True)
R = [smaller_step_R, normal_step_R, bigger_step_R]
batch_loss = loss_function(batch, dictionary, R, self.neuron_locations)
batch_loss.retain_grad()
batch_loss.backward()
batch_losses.append(batch_loss.item())


with
torch.no_grad():
self.smaller_step_loss += batch_losses[0]
self.normal_step_loss += batch_losses[1]
self.bigger_step_loss += batch_losses[2]


for
i, dictionary
in
enumerate(dictionaries):
dictionaries = self.get_spherical_step(dictionary, dictionary.grad, step_sizes)

self.smaller_step_dictionary, self.dictionary, self.bigger_step_dictionary = dictionaries


which is in turn called by the train_dictionary function:

def train_dictionary(self, training_batches, validation_set, num_epochs):
loss_function = LossFunction.LossFunction(self.penalty_type, self.lamb)
self.step_size = 0.1
self.dictionary.requires_grad_(True)

for epoch in range(num_epochs):
print(f"Starting epoch {epoch}")
training_batches = Preprocessing.shuffle_data(training_batches)

for batch_index, batch in enumerate(training_batches):
batch = batch.to(self.device)
if self.step_size < 1e-9:
self.dictionary.requires_grad_(False)
return

R = self.forward(batch)
self.optimizer_step(batch, loss_function)

if batch_index % 1000 == 0:
with torch.no_grad():
loss = loss_function(batch, self.dictionary, R, self.neuron_locations)
print(f"{batch_index}/{len(training_batches)} batches complete")
print(f"loss = {loss}")
print(f"current step size is: {self.step_size}")

with torch.no_grad():
_, acc, prec, recall = self.get_best_threshold(validation_set)

print(f"Epoch {epoch} complete. Accuracy, precision, and recall are as follows:\n{acc}\n{prec}\n{recall}")

self.dictionary.requires_grad_(False)

def optimizer_step(self, batch, loss_function):
if self.current_probe_step == self.max_probe_steps:
self.reset_probe()

self.current_probe_step += 1

with torch.no_grad():
smaller_step_R = torch.linalg.lstsq(self.smaller_step_dictionary, batch).solution
normal_step_R = torch.linalg.lstsq(self.dictionary, batch).solution
bigger_step_R = torch.linalg.lstsq(self.bigger_step_dictionary, batch).solution

dictionaries = [self.smaller_step_dictionary, self.dictionary, self.bigger_step_dictionary]
step_sizes = [self.step_size / 2, self.step_size, self.step_size * 2]

batch_losses = []
for i, dictionary in enumerate(dictionaries):
dictionary.requires_grad_(True)
R = [smaller_step_R, normal_step_R, bigger_step_R]
batch_loss = loss_function(batch, dictionary, R, self.neuron_locations)
batch_loss.retain_grad()
batch_loss.backward()
batch_losses.append(batch_loss.item())

with torch.no_grad():
self.smaller_step_loss += batch_losses[0]
self.normal_step_loss += batch_losses[1]
self.bigger_step_loss += batch_losses[2]

for i, dictionary in enumerate(dictionaries):
dictionaries = self.get_spherical_step(dictionary, dictionary.grad, step_sizes)

self.smaller_step_dictionary, self.dictionary, self.bigger_step_dictionary = dictionaries


I didn't use to have this error before, when I use a simple grid search for hyperparameter optimization. I only start to get this error when I tried using Optuna to do a Bayesian optimization. The error usually throws after I'm done with trial 0 and starts trial 1:

for target_dimension in range(upper_bound, lower_bound - 1, -1):

# Inner function to optimize lambda for a fixed target_dimension
def objective(trial):
nonlocal iteration

penalty_coefficient = trial.suggest_float("lambda", 1e-5, 10.0, log=True)

# Initialize model with pretrained dictionary if available
current_model = DictionaryLearning.DictionaryModel(
penalty_type=penalty_type,
penalty_multiplier=penalty_coefficient,
target_dimension=target_dimension,
original_dimension=original_dimension,
receptor_type=receptor_type,
neuron_locations=locations,
pretrained_dictionary=previous_dictionary,
is_random_init=is_random_init
).to(device)

# Train and evaluate model
current_model.train_dictionary(training_batches, validation_set, num_epochs=15)
cutoff, _, current_precision, current_recall = current_model.get_best_threshold(validation_set)

trial.set_user_attr("dictionary", current_model.dictionary)
trial.set_user_attr("model", current_model)
trial.set_user_attr("cutoff", cutoff)

current_stat_set = StatSet(space, penalty_coefficient, penalty_type, receptor_type, cutoff, current_model, validation_set)
current_f1_score = (2 * current_precision * current_recall) / (current_precision + current_recall)
sparsity_score = current_stat_set.average_utilization
locality_score = current_stat_set.interpretable_locality

lambdas.append(penalty_coefficient)
f1_scores.append(current_f1_score)
sparsity_scores.append(sparsity_score)
locality_scores.append(locality_score)

save_dictionary(save_path, iteration, current_model)
iteration += 1

# Return F1 score as the objective to maximize
return current_f1_score

# Run Bayesian Optimization on lambda for current target_dimension
study = optuna.create_study(direction="maximize")
study.optimize(objective, n_trials=20)

# Get the best F1 score and lambda for this target dimension
best_trial = study.best_trial
best_f1 = best_trial.value
best_lambda_for_dimension = best_trial.params["lambda"]

# Check if this target_dimension meets the F1 threshold
if best_f1 >= f1_threshold or first:
best_target_dimension = target_dimension
best_lambda = best_lambda_for_dimension
best_f1_score = best_f1

print(f"Best target_dimension: {best_target_dimension}, Best lambda: {best_lambda}, F1: {best_f1_score}")

best_dictionary = best_trial.user_attrs["dictionary"]
previous_dictionary = torch.clone(best_dictionary).to(device)

model = best_trial.user_attrs["model"]
cutoff = best_trial.user_attrs["cutoff"]

best_stat_set = StatSet(space, best_lambda, penalty_type, receptor_type, cutoff, model, validation_set)
best_stat_set.print_stats()
save_dictionary(save_path, "", model)

optimization_fig = plot_optimization_history(study)
slice_fig = plot_slice(study)

optimization_fig.figure.savefig("optimization_history.pdf", format="pdf")
slice_fig.figure.savefig("slice_plot.pdf", format="pdf")

if first:
first = False
else:
break


I looked this up on StackOverflow and included batch_loss.retain_grad() as a result, but the error is still there. In addition, why is the default only to calculate gradients for the leaf tensors? Any help would be really appreciated! Thank you.

Continue reading...