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    
相关推荐
LOVE️YOU几秒前
Python 在定义函数时,究竟定义的是什么?
python
段一凡-华北理工大学12 分钟前
大模型应用开发 100 天:Python + LLM 从入门到精通 day26~Prompt 调试与优化——A/B 测试与效果评估
windows·python·大模型·prompt·智能体·提示词工程·高炉智能化
Minecraft红客14 分钟前
以撒的结合
python·游戏·电脑·娱乐
Circ.25 分钟前
Python 调试 Litellm 本地大模型接口:模型列表与对话接口实操踩坑记录
开发语言·python
风早爽太1 小时前
用 ‌Render‌ 快速部署网站和常见问题解决
python·fastapi
海海不掉头发2 小时前
阶段五_实验34-38_学习总结
深度学习·神经网络·学习·机器学习
李航19832 小时前
给自己的图形引擎,配上了AI渲染,做设计真是太方便了
人工智能·python·计算机视觉·ai·ai编程
L@ncor2 小时前
第六章 框架开发实践 · 学习笔记(AutoGen / AgentScope / CAMEL / LangGraph)
python·框架·autogen·langgraph·agentscope
Zhou1411362 小时前
SpringSecurity_02_授权与高级功能
开发语言·python
布吉岛的石头3 小时前
Java 程序员第 49 阶段5:BERT 预训练目标 MLM+NSP 的工程含义
java·人工智能·深度学习·bert·transformer