Pytorch.device使用及理解

最新推荐文章于 2025-03-21 21:12:44 发布

介壳

最新推荐文章于 2025-03-21 21:12:44 发布

阅读量2.8w

点赞数 13

分类专栏： pytorch 文章标签： pytorch cuda

原文链接：https://blog.csdn.net/shaopeng568/article/details/95205345

版权

pytorch 专栏收录该内容

1 篇文章

订阅专栏

本文详细解析了PyTorch中torch.device的作用，包括如何选择设备（cuda或cpu），如何将模型迁移到指定设备，以及在多GPU环境下的处理方法。重点介绍了如何使用model.to(device)将模型与tensor移动到正确设备以优化GPU性能。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

Pytorch to（device）

Pytorch.device理解

device=cuda之类的
转载合并了一下其他同学的讲解.
链接: https://blog.csdn.net/shaopeng568/article/details/95205345
torch.device代表将torch.Tensor分配到的设备的对象。torch.device包含一个设备类型（‘cpu’或‘cuda’）和可选的设备序号。如果设备序号不存在，则为当前设备。如：torch.Tensor用设备构建‘cuda’的结果等同于‘cuda：X’，其中X是torch.cuda.current_device()的结果。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

这两行代码放在读取数据之前。

mytensor = my_tensor.to(device)

这行代码的意思是将所有最开始读取数据时的tensor变量copy一份到device所指定的GPU上去，之后的运算都在GPU上进行。

这句话需要写的次数等于需要保存GPU上的tensor变量的个数；一般情况下这些tensor变量都是最开始读数据时的tensor变量，后面衍生的变量自然也都在GPU上.

如果是多个GPU
在代码中的使用方法为：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = Model()
if torch.cuda.device_count() > 1:
  model = nn.DataParallel(model，device_ids=[0,1,2])

pytorch中model=model.to(device)用法

https://blog.csdn.net/weixin_36670529/article/details/106068761
https://blog.csdn.net/weixin_36670529/article/details/106068761
device=torch.device(“cpu”)代表的使用cpu，而device=torch.device(“cuda”)则代表的使用GPU。
当我们指定了设备之后，就需要将模型加载到相应设备中，此时需要使用model=model.to(device)，将模型加载到相应的设备中。将由GPU保存的模型加载到CPU上。