机器学习:通过 to 方法调用GPU/CPU

之前在自己写代码时候费尽周折也并没有成功调用CPU。
今天在学长给的代码里面看到了

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print('device:', device)

这次竟然成功调用了独显。关于环境,应该是之前配置上成功了一部分,就是安装CUDA和cudnn以及tensorflow-gpu(tensor可以在GPU上运行,Numpy不可.
现在来研究一下学长是怎样实现的。
追踪device发现在加载数据调用网络训练时处出现了 .to(device)

    net = ResNet50(num_classes=CLASS_NUMBER).to(device)

在GPU上进行运算

	inputs, labels = inputs.to(device), labels.to(device)

把数据导入到GPU上去.
to( ) 可用于容易地将对象移动到不同的设备(代替以前的cpu()或cuda()方法)
从网上找了一种使用模式:就是先设定device,导入数据到GPU,再设定计算位置在GPU上

# 开始脚本,创建一个张量
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

...

# 但是无论你获得一个新的Tensor或者Module
# 如果他们已经在目标设备上则不会执行复制操作
input = data.to(device)
model = MyModule(...).to(device)

参考:https://ptorch.com/news/189.html

你可能感兴趣的:(深度学习,tensorflow)