PART 1: 使用GPU运行程序的方法

a 首先确定计算机或者服务器有可以使用的GPU

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device

注意:
a) 只有一块GPU则使用:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

b) 有多块GPU则使用:

device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")

cuda:1----定义使用GPU1,GPU的数量通常从0开始排序
c) 获取能够使用的GPU数量

`torch.cuda.device_count()`

b 将网络转为cuda模式

net = Net().to(device)

注意:Net() 是已经定义好的网络类,用Net()来初始化net

c 将数据转为cuda模式

data.to(device)                #将数据转化为使用指定的GPU可计算形式0
data.cpu()                       #将GPU可计算的形式的数据转化为CPU的可计算形式

注意:
在使用GPU进行计算时,要搞清楚何时使用什么样式的数据,在计算时才不会出错

PART 2: CPU使用利用GPU训练好的模型

a 首先定义CPU计算

device = torch.device('cpu')

b 定义一个相同的在CPU情况下的网络

net = Net()

c 导入网络参数,并指定参数map_location=device,即网络的参数映射到基于CPU计算网络

net.load_state_dict(torch.load("(已经训练好并且保存下来的模型地址).pkl", map_location=device))

下面就可以在CPU上直接使用GPU训练好的模型了!

PART 3: 几个小技巧

1)数据加载

data_loader = Data.DataLoader(dataset, batch_size=4, shuffle=True,
                              num_workers=8,pin_memory=True)

pin_memory=True,这里的参数将数据定义为锁页内存,这样转移到GPU上的速度会更快一点

2)需要什么在GPU上运行,就将什么内容传到GPU上。

.to(device)

PyTorch初学者,请多指教!!!

Logo

为开发者提供学习成长、分享交流、生态实践、资源工具等服务,帮助开发者快速成长。

更多推荐