首先需要确保系统上安装了CUDA支持的NVIDIA GPU和相应的驱动程序。
基本步骤如下
检查CUDA是否可用：
使用 torch.cuda.is_available() 来检查CUDA是否可用。
指定设备：
可以使用 torch.device(“cuda:0”) 来指定要使用的GPU。如果系统有多个GPU，可以通过改变索引（例如 “cuda:1”）来选择不同的GPU。
将模型移至GPU：
使用 .to(device) 方法将模型移至GPU。例如，如果 model 是PyTorch模型，且 device 是之前定义的设备，使用 model.to(device)。
将数据移至GPU：
同样地，将数据（例如张量）也需要移至GPU。例如，如果 data 是一个张量，使用 data = data.to(device)。
进行计算：
然后可以像平常一样进行模型的训练或推断，PyTorch会在GPU上执行计算。
将数据移回CPU（如果需要）：
如果你需要将结果移回CPU，可以使用 .to(“cpu”)。
简单的示例代码如下：
import torch
import torch.nn as nn
# 检查CUDA是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 创建一个模型示例
model = nn.Linear(10, 5).to(device)
# 创建一些数据并移至GPU
input_data = torch.randn(10).to(device)
# 在GPU上进行前向传播
output = model(input_data)
# 将输出移回CPU（如果需要）
output = output.to("cpu")
在使用GPU进行计算时，要管理好内存，因为GPU内存通常比系统内存有限。如果遇到内存溢出错误，可能需要减少批量大小或优化模型。
原文地址:https://blog.csdn.net/qq_49370210/article/details/135984165
本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。
如若转载，请注明出处：http://www.7code.cn/show_64445.html
如若内容造成侵权/违法违规/事实不符，请联系代码007邮箱：suwngjj01@126.com进行投诉反馈，一经查实，立即删除！
