Pytorch 容器 - 5. Module中的参数转换/转移:cpu(), cuda(), float(), double(), half()

本文介绍了PyTorch中如何在CPU和GPU间转换模型参数,包括model.cpu()、model.cuda()以及float()、half()、double()等函数的使用,强调了在GPU训练中需手动调用cuda()并将模型、数据和损失函数一并转移。
摘要由CSDN通过智能技术生成

1. cpu()

对于模型的cpu()函数,比如 model.cpu(),只是把 parameters 和 buffers 进行转移

2. cuda(device=None)

在PyTorch中,算法优先使用CPU,不会自动使用GPU,因此需要在程序中显示指定。如果模型要在gpu上进行训练,需要再将模型参数放到优化器之前调用cuda() 将其加载到GPU。

  • 在网络模型的训练过程中,模型,数据(input和gt)和损失函数都需要转移到cuda
  • 使用 torch.cuda.is_available() 判断gpu是否可用
  • 除了 cuda(device) 函数外,还可以使用 to(device) 来进行设备转换。其中 to(device) 中的参数是必须的,cuda()中的参数是可选的。对于单机单卡(单GPU)的模型训练,代码如下:
model = Model()
# 第一种定义 device的方式
device = torch.device("cuda:1") 
# 第二种定义 device的方式
device = torch.dev
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值