如何解决pytorch中的“OutOfMemoryError：CUDA内存不足”？

Question

在 Colab 中，我使用经过训练的模型预测 2448x2448 的数组，其中有 7 个类（输入= (2448, 2448, 3) 和输出= (2448, 2448, 7)。

for idx in range(len(test_dataset)):

image, gt_mask = test_dataset[idx]
image_vis = test_dataset_vis[idx][0].astype('uint8')
x_tensor = torch.from_numpy(image).to(DEVICE).unsqueeze(0)
# Predict test image
pred_mask = best_model(x_tensor)
pred_mask = pred_mask.detach().squeeze().cpu().numpy()
# Convert pred_mask from `CHW` format to `HWC` format
pred_mask = np.transpose(pred_mask,(1,2,0))
# Get prediction channel corresponding to foreground
pred_urban_land_heatmap = pred_mask[:,:,select_classes.index('urban_land')]
pred_mask = colour_code_segmentation(reverse_one_hot(pred_mask), select_class_rgb_values)
# Convert gt_mask from `CHW` format to `HWC` format
gt_mask = np.transpose(gt_mask,(1,2,0))
gt_mask = colour_code_segmentation(reverse_one_hot(gt_mask), select_class_rgb_values)
cv2.imwrite(os.path.join(sample_preds_folder, f"sample_pred_{idx}.png"), np.hstack([image_vis, gt_mask, pred_mask])[:,:,::-1])

visualize(
    original_image = image_vis,
    ground_truth_mask = gt_mask,
    predicted_mask = pred_mask,
    pred_urban_land_heatmap = pred_urban_land_heatmap
)

但是我明白了

OutOfMemoryError：CUDA 内存不足。尝试分配 366.00 MiB。 GPU 0 的总容量为 14.75 GiB，其中 357.06 MiB 可用。进程 224843 有 14.40 GiB 内存正在使用。在已分配的内存中，13.94 GiB 由 PyTorch 分配，344.24 MiB 由 PyTorch 保留但未分配。如果保留但未分配的内存很大，请尝试设置 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True 以避免碎片。请参阅内存管理文档 (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

我在 colab 中有 15 GB 的 GPU，但是当我到达这条线时

pred_mask = best_model(x_tensor)

分配的 GPU 内存在分配图中达到顶部。

Answer 1

您的模型太大或您的输入太大。你没有太多选择。使用较小的模型或使用较小的输入。对于大多数网络来说，2448x2448x3 通常是一个非常大的数组。如果您使用图像，他们通常会将 224x224 或 512x512 等图像作为输入，因此您需要调整大小或进行平铺。