PyTorch在Windows下按预期运行，但在Ubuntu下的较大图像上失败

Question

我训练过一个用于图像分割的segmentation_models_pytorch.PSPNet模型。为了进行预测，我在PyTorch tensor中加载整个图像，并使用384x384像素窗口对其进行扫描。

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

我的Windows机器有6Gb的GPU，而Ubuntu有8GB的GPU。当加载所有模型时，它们消耗了大约1.4GB的GPU。在Windows上处理大图像时，内存消耗增加到1.7GB GPU。

在Windows环境下，该模型可以处理25M像素图像。在Ubuntu下，相同的代码只能处理高达5M像素的图像。调试很困难，因为我只有ssh访问Ubuntu机器。是什么导致了这一差异，以及如何调试此问题？

Answer 1

你可以检查库达版本，火炬版本。

要进行调试，您可以在image_tensor.shape之前打印model.predict，也许您正在linux上运行更大的批处理大小。

Answer 2

第一种解决办法：

我首先应该检查CUDA、Pytorch和其他一般python库的库版本。如果它们都是相等的，那么检查在Ubunto和您的GPU驱动器/CUDA中的库的兼容性。

正如Jimmy所评论的，您可以在Python上直接使用conda环境或虚拟环境来控制库的版本。

调试

您可以安装允许调试的Python。如果不可能，一种非常简单但效率低下的方法是在代码的许多部分上打印，以查看什么与Windows上的执行完全不同。