GPU

话不多说，直接进入主题。

✔ 判断能不能使用GPU
可能有多种原因会导致不能使用GPU，比如PyTorch安装的是CPU版的，显卡驱动没有正确安装等。下面的 if 语句在正常的情况下会返回 True：

if torch.cuda.is_available():
    print('PyTorch can use GPU on current machine!')

文章来源：https://www.codelast.com/
✔ 设置模型使用GPU

model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(your_model_file_path))
model.eval()  # 设置成evaluation模式
if torch.cuda.is_available():
    print('PyTorch can use GPU on current machine!')
    device = torch.device("cuda")
    model.to(device)

your_model_file_path 是模型文件的路径。

阅读更多