Transformer Model can predict values in training but not on evaluation

Deadrosas: <p>I am creating a vision transformer for visual odometry. The model I have now is based on the ViViT which is a transformer capable of establishing spatial and temporal relationships.</p>
<p>In training the model is capable of predicting something  ( Top image) But when i change the model to eval mode and try to predict on the same data or other similar I only get a straight line (Bottom Image).</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/9/291efba6086d8ae948e0c86bb09045c72b6b3eed.png" data-download-href="https://discuss.pytorch.org/uploads/default/291efba6086d8ae948e0c86bb09045c72b6b3eed" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/9/291efba6086d8ae948e0c86bb09045c72b6b3eed_2_324x499.png" alt="image" data-base62-sha1="5RLXkXgGwQbqX3h4Q5guzpchwHH" width="324" height="499" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/9/291efba6086d8ae948e0c86bb09045c72b6b3eed_2_324x499.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/9/291efba6086d8ae948e0c86bb09045c72b6b3eed_2_486x748.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/2/9/291efba6086d8ae948e0c86bb09045c72b6b3eed.png 2x" data-dominant-color="F7F7F7"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">603×930 42.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>In blue the Ground Truth and in orange the predictions</p>
<p>I stack 2 images on top of each other and feed them to the transformer as seen in the training and test functions</p>
<p>An image as the following format →  torch.rand(1, 1, 2, 370, 1241) ↔ (Batches, Channels, NumOfFrames, Height, Width)</p>
<pre><code class="lang-auto">for epoch in range(epochs):

    ### Training 
    train_loss = 0.0
    running_loss = 0.0
    loss_values = []
    pred_poses = np.empty([len(train_dataset), 3])
    pred_poses[0] =  train_dataset.getPoses()[0][:, 3:][:3].flatten()

    for i in tqdm(range(0, len(train_dataset) - 1)):
        
        imgT = train_dataset[i]["image"].unsqueeze(1)
        imgT2 = train_dataset[i + 1]["image"].unsqueeze(1)

        imgMerge = torch.cat((imgT, imgT2), dim=1).unsqueeze(0)
        optimizer.zero_grad()

        with torch.set_grad_enabled(True):
            outputs = model(imgMerge) #.cuda()
            print("Preds: ", outputs.flatten())

            labels = train_dataset[i + 1]["pose"] - train_dataset[i]["pose"]
            print("GT Pose: ", labels[:, 3:][:3].flatten())
            pred_poses[i + 1] = (np.add(pred_poses[i], outputs[:, :3].cpu().flatten().detach().numpy()))
            loss = criterion(outputs, labels)
            train_loss += loss
            # print(loss)
            loss.backward()
            optimizer.step()
            loss_values.append(loss.item())

        train_loss /= len(train_dataset)

        print(f"Training Loss: {train_loss:.4f}")

    plot_loss(loss_values)
    plot_odometry(train_dataset, pred_poses)

    ### Testing

    test_loss = 0.0
    
    running_corrects = 0
    test_pred_poses = np.empty([len(test_dataset), 3])
    test_pred_poses[0] =  test_dataset.getPoses()[0][:, 3:][:3].flatten()

    for i in tqdm(range(0, len(test_dataset) - 1)):
        imgT = test_dataset[i]["image"].unsqueeze(1)
        imgT2 = test_dataset[i + 1]["image"].unsqueeze(1)

        imgMerge = torch.cat((imgT, imgT2), dim=1).unsqueeze(0)

        # print(imgMerge.shape)

        with torch.inference_mode():
            outputs = model(imgMerge) #.cuda()
            # print(pred_poses)
            print(outputs.flatten())
            test_pred_poses[i + 1] = (np.add(test_pred_poses[i], outputs[:, :3].cpu().flatten().numpy()))
            print("GT: {} | Pred: {}".format(np.array(test_dataset[i + 1]["pose"])[:, 3:][:3].flatten(), test_pred_poses[i + 1]))
            # loss = criterion(outputs, dataset[i + 1]["pose"] - dataset[i]["pose"])
            # test_loss += loss


        # test_loss /= len(dataset)

        # running_loss += loss.item() * (poses[i + 1] - poses[i]).size(0)

    epoch_loss = running_loss / len(test_dataset)
    plot_odometry(test_dataset, test_pred_poses)
</code></pre>
<p>And my custom loss function:</p>
<pre><code class="lang-auto">def loss_fn(outputs, labels):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    yaw=math.atan2(labels[1][0],labels[0][0])
    pitch=math.atan2(-labels[2][0],math.sqrt(labels[2][1]*labels[2][1]+labels[2][2]*labels[2][2]))
    roll=math.atan2(labels[2][1],labels[2][2])

    translation_loss = torch.nn.functional.mse_loss(labels[:, 3:][:3].to(device).flatten().float(), outputs[:, :3].flatten().float())
    rotation_loss = torch.nn.functional.mse_loss(torch.tensor([yaw, pitch, roll]).to(device).float(), outputs[:, 3:].flatten().float())

    return translation_loss + rotation_loss * 100 
</code></pre>
<p>For more context, In eval mode the transformer seems to predict roughly the same values for every single input which doesn’t make much sense to me. Even when the transformer is not trained the output is a straight line. In training this doesn’t happen at all. I also used the architecture I’ve built for video classification and it performed as expected. So I assume the architecture is not the problem…</p>
<p>I can’t figure what is happening and any help is appreciated.</p>