Tricky forward pass pytorch

user_3791408 0
Viewed 298

I have a precipitation map timeseries dataset with input shape (None, seq_length =7, c = 75, w=112, h=112) output shape (None, lead_times = 60, c=51, w=28, h=28). The model (Conv downsampler + ConvGRU + Axial Attention) predicts precipitation in a 28x28 region in the middle with 51 categorical precipitation intervals and is conditioned with 60 different lead times (5, 10, ..., 300 minutes).

Right now my forward pass looks like this:

    def forward(self, imgs):
      """It takes a rank 5 tensor
      - imgs [bs, seq_len, channels, h, w]
      """

      # Compute all timesteps, probably can be parallelized
      res = []
      for i in range(self.forecast_steps):
          x_i = self.encode_timestep(imgs, i)

          out = self.head(x_i)
          res.append(out)
      res = torch.stack(res, dim=1)
      return res

Here imgs is the input tensor without lead time encoding, so only 15 channels. The imgs is then one-hot encoded for each respective lead time and the output is the entire predicted time series (5-300min). However this leads to severe memory issues even with batch_size = 1 so I want the forward loop to only do one random lead time at a time. I am training this with pytorch-lightning module for easier parallelization so I don't have much control of the training loop.

The issue is that the effective batch size with this training loop is 60*batch_size. The paper solves this by only doing one random lead time per sample, which now makes sense to me. This solves the memory issue by allowing effective minimum batch size to be 1. How can pass a random integer (the lead time) to the forward pass and couple it with the correct Y when pytorch-lightning computes the loss?

I want

y_hat = forward(self, X[n], lead_time=random)
...
loss(y_hat-Y[n,lead_time,:,:])

My code is available at https://github.com/ValterFallenius/metnet.

weather pytorch-lightning pytorch python
1 Answers

I figured out how to fix it. Once explaining the problem to someone else, I realized how simple the solution was...

def forward(self, imgs,lead_time):
    """It takes a rank 5 tensor
    - imgs [bs, seq_len, channels, h, w]
    - lead_time #random int between 0,self.forecast_steps
    """

    x_i = self.encode_timestep(imgs, lead_time)

    out = self.head(x_i)
    res.append(out)
    res = torch.stack(res, dim=1)
    return res

The trick was to simply add the lead_time variable to the training_step method:

def training_step(self, batch, batch_idx):
    x, y = batch


    lead_time = np.random.randint(0,self.forecast_steps)

    y_hat = self(x.float(),lead_time)
    loss = F.mse_loss(y_hat, y[:,lead_time])
    pbar = {"training_loss": loss}
    return {"loss": loss, "progress_bar": pbar}
user_37914080
Related