Stable Baselines3：如何为model.predict指定特定GPU

Stable-Baselines3（SB3）基于PyTorch实现，给model.predict指定GPU本质是控制PyTorch的设备分配，以下是两种常用方法：

方法1：通过环境变量指定GPU #

在运行测试脚本前，设置CUDA_VISIBLE_DEVICES环境变量，限制脚本仅能访问指定GPU，模型的所有操作（包括predict）都会自动使用该GPU：

# 示例：仅让脚本使用编号为1的GPU（GPU编号从0开始）
CUDA_VISIBLE_DEVICES=1 python your_test_script.py

方法2：在代码中手动指定设备 #

直接在代码里定义目标GPU设备，将模型加载或迁移到该设备上，确保predict调用时使用指定GPU：

1. 导入PyTorch并定义目标设备：

import torch
# 指定使用编号为1的GPU，若要使用CPU则写"cpu"
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")

2. 加载模型时指定设备：

from stable_baselines3 import PPO

# 加载模型时直接指定设备
model = PPO.load("your_trained_model.zip", device=device)

3. 若模型已加载到其他设备，可手动迁移：

# 将模型的策略网络迁移到目标设备
model.policy.to(device)

完成以上设置后，调用model.predict()时就会自动使用你指定的GPU执行推理。

内容的提问来源于stack exchange，提问作者MarcoM