[Python] How do I avoid CPU for loops when inputs are a function of previous outputs?

I have a simple reinforcement model in pytorch that makes decisions about where it wants to be and receives its current position (e.g. from a GPS) as an input. For the sake of this example let's just assume it can just move anywhere at any point (ie. it can just specify its position). In the real example we'd either output a movement or the activation would augment the desired position to what is achievable.

class Traveller(torch.nn.Module):
def __init__(
self,
num_inputs: int,
hidden_size: int,
num_layers: int,
dropout: float,
):
super(DeepStatefulMVO, self).__init__()

# The first LSTM layer
self.lstm_layers = torch.nn.LSTM(
input_size=num_inputs + 2,
hidden_size=hidden_size,
num_layers=num_layers,
dropout=dropout,
batch_first=True
)

self.output_layer = torch.nn.Linear(
in_features=hidden_size,
out_features=2
)

self.activation = torch.nn.Tanh()

def forward(
self,
inputs: torch.Tensor, # this is some independent data that is known ahead of time, e.g. a clock, the wind, lighting conditions, etc.
) -> torch.Tensor:

assert (inputs.shape[1] == 1), "window should always be 1 since we're not using sequences"

hidden_state: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
prev_position = torch.zeros(
(1, 1, 2), device=inputs.device, dtype=inputs.dtype
)
positions_list = []
for t in range(inputs.size(0)):
current_inputs = inputs[t : t + 1]
input_data = torch.cat(
[
prev_position,
current_inputs
],
dim=2,
)
lstm_out, hidden_state = self.lstm_layers(input_data, hidden_state)
out = self.output_layer(lstm_out)
position = self.activation(out)
prev_position = position
positions_list.append(position)

positions = torch.cat(positions_list, dim=0).squeeze(dim=1)
assert isinstance(positions, torch.Tensor)
assert positions.shape == (inputs.shape[0], inputs.shape[2])
return positions


The output is a set of positions that are fed into a utility function which computes some loss (with a gradient).

This code works fine, but the issue is that the for loop is orchestrated from the CPU. Previously (before I had the positions vector being used as an input), I could just pass in an entire batch of data with a large sequence dimension, and everything would be computed in one step, but the LSTM layers wouldn't know their current position, so the model was unable to optimise the loss effectively. Moving to this approach improved training convergence drastically but reduced GPU utilisation from about 95% to 20%.

How can I retain this functionality where the input is a function of the previous output, without resorting to CPU-bound for-loops?

Continue reading...