PyTorch下,使用list放置模块,导致计算设备不一的报错

报错

在复现 Transformer 代码的训练阶段时,发生报错:

bash 复制代码
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

解决方案

通过next(linear.parameters()).device确定 model 已经在 cuda:0 上了,同时输入 model.forward()的张量也位于 cuda:0。输入的张量没什么好推敲的,于是考虑到模型具有多层结构,遂输出每层结构的设备信息,model.encoder -> model.encoder.sublayer[0] ··· ···

测试发现,model.encoder.sublayer[0] 之后的模块的设备信息均位于 cpu,原因是构造这部分模块时,由于需要多个相同的模块,使用了 list 来存放模块:

python 复制代码
# module: 需要深拷贝的模块
# n: 拷贝的次数
# return: 深拷贝后的模块列表
def clones(module, n: int) -> list:
    return [copy.deepcopy(module) for _ in range(n)]

显然 list 不支持 GPU,需要用 PyTorch 提供的代替:

python 复制代码
def clones(module, n: int):
    return nn.ModuleList([copy.deepcopy(module) for _ in range(n)])

ModuleList 把子模块存入列表,能像 Python 里普通的列表被索引,最重要的是能使内部的模块被正确注册,并对所有的 Module 方法可见。[Source](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html#torch.nn.ModuleList)

成功解决!

相关环境

bash 复制代码
python                    3.11.7               he1021f5_0
pytorch                   2.1.2           py3.11_cuda12.1_cudnn8_0    
相关推荐
Python自动化直播5 分钟前
智能脚本批量管理拆开看:导入、导出、重写、试听
前端·python·ai·秒播
AI人工智能+25 分钟前
机动车登记证识别技术通过“图像矫正+版面分析+高精度OCR+业务校验”四步流水线,实现对“大绿本”的秒级精准识别
人工智能·深度学习·ocr·机动车登记证识别
具身AGI1 小时前
三维动态只当训练信号,不进推理
人工智能·深度学习·计算机视觉
乐迪信息1 小时前
AI算法盒子船体烟火检测:港口火灾早预警
人工智能·深度学习·算法
拜小白1 小时前
【论文解读】用多标签图割从多视角全景图精确估计房间布局
深度学习·计算机视觉·全景图
承渊政道2 小时前
HIL-UMI论文解读:用手持示范补齐机器人操作短板——对启元Q1机器人开发的启发
人工智能·深度学习·机器人·论文解读·具身智能·启元q1
伤心的梧桐叶3 小时前
Picset与创客贴比等待时间,没做完的任务不能从表里消失
python
没有气的汽水3 小时前
一次 SQS 消费者静默退出的排查:问题最后竟然出在日志里
python
用户450030140233 小时前
物业催缴系统为什么总在"冷却"这一步翻车:账期建模与规则命中的几个取舍
python
Web极客码3 小时前
决策模型 + LLM:用 Python 拆解浏览器自动化工作流
java·python·自动化