[SOLVED] PyTorch runs as expected under Windows, but fails on larger images under Ubuntu

Issue

I've trained a segmentation_models_pytorch.PSPNet model for image segmentation. For prediction I load whole image in PyTorch tensor and scan it with 384x384 pixels window.

result = model.predict(image_tensor[:, :, y:y+384, x:x+384])

My Windows machine has 6Gb GPU, while Ubuntu has 8 Gb GPU. When all models are loaded they consume some 1.4 Gb GPU. When processing a large image on Windows the memory consumption increases to 1.7 Gb GPU.

Under Windows the model can handle 25 M pixel images. Under Ubuntu the same code can only process up to 5 M pixel image. Debugging is difficult because I only have ssh access to the Ubuntu machine. What could cause this discrepancy and how to debug this issue?

Solution

You can check cuda version, torch version.

To debug, you can just print image_tensor.shape before model.predict, maybe you are running a larger batch size on linux machine.

Answered By - Jimmy

Answer Checked By - Katrina (WPSolving Volunteer)