PyTorch中使用了张量类型,而不用numpy的array,就是为了可以在GPU上运行代码,那我们怎么样才能使用GPUs来加速运行呢。其实非常简单,几条语句就可以完成了。
1. 基本语句
- 模型放到一个GPU上运行
model.gpu()
- 将张量放到GPU上
mytensor = my_tensor.gpu()
注: 仅仅调用my_tensor.gpu()
函数并不会将张量复制到GPU上,所以一定要指定一个新的张量mytensor
,然后在 GPU上用这个新的张量。- 模型放在多个GPU上运行
上文中的model.gpu()
默认只使用一个GPU,如果你有多个GPU的话,应该这样子调用model = nn.DataParallel(model)
2. 代码示例
在这个例子中,模型获取一个输入,经过线性变换后得到输出结果,具体的看下面的代码吧~
- 包导入及参数定义
import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader
# Parameters and DataLoaders
input_size =