判断能不能使用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 是模型文件的路径。

 inference/predict的时候使用GPU
对一次 inference 来说，假设模型的输入数据为 model_input_tensor（torch.Tensor类型），那么计算模型输出的方法是：

if torch.cuda.is_available():  # GPU available
    model_input_tensor = model_input_tensor.to(torch.device('cuda'))
model_output = model(model_input_tensor)  # inference

 检查程序跑起来之后是不是真的用了GPU
用 nvidia-smi 命令来查看。

文章来源：https://www.codelast.com/
➤➤ 版权声明 ➤➤ 
转载需注明出处：codelast.com 
感谢关注我的微信公众号（微信扫一扫）：

torch.save(the_model.state_dict(), PATH)

对应地，加载模型这样做：

the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))

另一种方法是不推荐的：

torch.save(the_model, PATH)

对应地，加载模型这样做：

the_model = torch.load(PATH)

文章来源：https://www.codelast.com/
这两者的区别：第1种方法只保存了模型的参数，而第2种方法保存了整个模型(结构+参数)，所以第2种方法保存出来的文件体积会比第1种方法大。
使用第2种方法的话，序列化的数据将绑定到所使用的特定类和确切的目录结构，因此在其他项目中使用或经过一些大的重构后，它可能会失效。
文章来源：https://www.codelast.com/
➤➤ 版权声明 ➤➤ 
转载需注明出处：codelast.com 
感谢关注我的微信公众号（微信扫一扫）：