use index assignment instead of torch.where to improve memory ussage in simulated annealing and greedy

[Laouen]

The current simulated annealing and greedy algorith uses torch.where to assign the new values to the current solutions. This creates an intermediate new tensor that consumes memory. Instead the new simulated annealing uses index assignment which updates in place the values of the current solution tensors and consumes less memory as explained in the following example:

To understand the memory consumption differences between the two PyTorch operations, let's break down what each operation does:

1. Using torch.where

best_solution = torch.where(
    new_global_maxima.unsqueeze(1).expand((-1, order)),  # match correct shape
    new_solution,
    best_solution
)

• torch.where:
  • This function returns a tensor composed of elements from new_solution where the condition (i.e., new_global_maxima.unsqueeze(1).expand((-1, order))) is True and from best_solution where it is False.
  • Memory Consumption:
    • The torch.where operation generally creates a new tensor to hold the result. This new tensor will have the same shape as best_solution and new_solution.
    • The intermediate tensor created by new_global_maxima.unsqueeze(1).expand((-1, order)) also consumes memory temporarily, which is shaped according to the expanded size.
    • Thus, memory is required for the output tensor and the temporary expanded tensor.

2. Index assignment

best_solution[new_global_maxima] = current_solution[new_global_maxima]

• Index-based Assignment:
  • This operation directly updates the best_solution tensor at the indices specified by new_global_maxima with values from current_solution at the same indices.
  • Memory Consumption:
    • There is no new tensor creation here; instead, best_solution is updated in place.
    • Memory is used only for the subset of elements accessed and copied from current_solution to best_solution, not for creating a new tensor.
    • This operation is generally more memory-efficient since it avoids the overhead of creating a new tensor.

Summary of Memory Consumption

• torch.where: Consumes more memory due to the creation of a new tensor and potentially larger intermediate tensors (e.g., from expanding).
• Index Assignment: More memory-efficient as it updates tensors in place without needing to allocate memory for a new tensor.

When to Use Which

• Use torch.where when you need a new tensor and cannot modify the original best_solution.
• Use index assignment when you want to update an existing tensor in place and save memory.

In general, if you are concerned about memory usage, prefer the index assignment method as it avoids the overhead of allocating new memory for the result.

The idea is to always use index assignment given that we are overriding the value of the current tensors and not creating new ones.