GPU能大幅提升计算速度,为了达到GPU加速训练的目的,就要将进行运算的数据传入GPU,回想一下神经网络,所有的数据不外乎是模型参数和数据集,所以我们只要将模型参数和数据集传入GPU即可。
用 .cuda() 可以达到这个目的
class Net(torch.nn.Module):
def __init__(self):... #我只是省略了网络详细内容 0.0
def forward(self, x):...
return x
my_cnn=Net().cuda() #将网络parameters传入GPU
数据集同理,读取后后边加上 .cuda() 即可
注意GPU里的数据和CPU里的是不能混用, 需要将数据转成CPU的话,后边加上 .cpu() 即可。
附上几个常用的GPU指令:
torch.cuda.is_available()
cuda是否可用;
torch.cuda.device_count()
返回GPU数量;
torch.cuda.get_device_name(0)
返回GPU名字,设备索引默认从0开始;
torch.cuda.current_device()
返回当前设备索引
torch.cuda.empty_cache()
释放缓存分配器当前持有的所有未占用的缓存显存
了解更多:https://pytorch.apachecn.org/docs/1.2/cuda.html