请使用 torch.load 和 map_location=torch.device('cpu') 到 map 你的存储到 CPU

Question

# set the computation device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Load model checkpoint
checkpoint = 'checkpoints/checkpoint_ssd300.pth.tar'
checkpoint = torch.load(checkpoint)
start_epoch = checkpoint['epoch'] + 1
print('\nLoaded checkpoint from epoch %d.\n' % start_epoch)
model = checkpoint['model']
model = model.to(device)
model.eval()

当我尝试运行此代码块时，出现以下问题：

RuntimeError：尝试在 CUDA 设备上反序列化 object，但 torch.cuda.is_available() 为 False。 如果你在一台只有 CPU 的机器上运行，请使用 torch.load 和 map_location=torch.device('cpu') 到 map 你的存储到 CPU。

Answer 1

错误消息表明您正在尝试加载在 GPU（CUDA 设备）上训练的 model 检查点，但您当前的机器没有 GPU 或 CUDA 不可用。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 试图确定 CUDA 在当前机器上是否可用，如果可用，则将设备变量设置为 'cuda '，否则将其设置为'cpu'。

行 checkpoint = torch.load(checkpoint) 试图从指定文件加载 model 检查点，但它试图在导致错误的“cuda”设备上这样做。

要解决此问题，您可以使用 torch.load function 的 map_location 参数来指定应将 model 加载到“cpu”设备上，而不是“cuda”设备上。

checkpoint = torch.load(checkpoint, map_location=torch.device('cpu'))

这样 model 将被加载到 CPU 设备上，即使 CUDA 设备被用来训练它。