Pytorch使用GPU

Posted 2023-01-12 汪小敏同学

Pytorch使用GPU

pytorch默认是不使用gpu的，即使我们已经安装好了支持的cuda版本和cudnn，所以要使用gpu需要在程序里面设置一下。

步骤

import torch
use_gpu = torch.cuda.is_available()

定义一个检查GPU是否可用的全局变量

加载模型、并设置模式：

    #下载预训练模型
    model = torchvision.models.segmentation.deeplabv3_mobilenet_v3_large(
        pretrained=True
    )
    if use_gpu:
        model.cuda()

推理和训练一样。

设置输入张量为cuda模式

 if use_gpu:
     tensor_cv = tensor_cv.cuda()

只需要设置第一个张量就可以了，后面关于该张量的所有计算都会在GPU上运行。

从tensor到numpy

result_person = result_person.cpu().detach().numpy()

注意需要使用.cpu()方法来把张量从gpu中取出来。