Pytorch如何将嵌套的dict类型数据加载到GPU

在PyTorch中,您可以使用.to(device)方法将嵌套的字典中的所有支持的Tensor对象转移到GPU。以下是一个简单的例子

python 复制代码
import torch
 
# 假设您已经有了一个名为device的GPU设备对象
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
# 嵌套的字典,其中包含一些Tensors
nested_dict = {
    'a': torch.randn(2, 2),
    'b': {
        'b1': torch.randn(2, 2),
        'b2': torch.randn(2, 2)
    },
    'c': torch.randn(2, 2)
}
 
# 将嵌套字典中的所有Tensors移动到GPU
def to_gpu(data):
    if isinstance(data, dict):
        return {k: to_gpu(v) for k, v in data.items()}
    elif isinstance(data, list):
        return [to_gpu(i) for i in data]
    elif isinstance(data, tuple):
        return tuple([to_gpu(i) for i in data])
    elif torch.is_tensor(data) and data.device != device:
        return data.to(device)
    else:
        return data
 
nested_dict_gpu = to_gpu(nested_dict)
 
# 检查是否所有Tensors都已移动到GPU
for k, v in nested_dict_gpu.items():
    if torch.is_tensor(v):
        assert v.device == device

这个函数to_gpu会递归地检查字典中的每个元素,如果是Tensor类型并且不在GPU上,就会使用.to(device)方法转移它。您需要先设置device变量指向您的GPU设备。如果没有GPU可用,它会默认使用CPU。

相关推荐
啊文师兄13 分钟前
使用 Pytorch 搭建视频车流量检测资源(基于YOLO)
人工智能·pytorch·yolo
使者大牙24 分钟前
【LLM学习笔记】第三篇:模型微调及LoRA介绍(附PyTorch实例)
人工智能·pytorch·python·深度学习
scdifsn28 分钟前
动手学深度学习10.1. 注意力提示-笔记&练习(PyTorch)
pytorch·笔记·深度学习·注意力机制·注意力提示
陌上阳光32 分钟前
动手学深度学习69 BERT预训练
人工智能·深度学习·bert
新手小白勇闯新世界2 小时前
点云论文阅读-1-pointnet++
论文阅读·人工智能·深度学习·神经网络·计算机视觉
小菜日记^_^2 小时前
BEAGLE: Forensics of Deep Learning Backdoor Attack for Better Defense(论文阅读)
论文阅读·人工智能·深度学习·sp·ai安全·backdoor 后门攻击·安全四大
小毕超2 小时前
基于 PyTorch 从零手搓一个GPT Transformer 对话大模型
pytorch·gpt·transformer
千天夜2 小时前
激活函数解析:神经网络背后的“驱动力”
人工智能·深度学习·神经网络
m0_523674212 小时前
技术前沿:从强化学习到Prompt Engineering,业务流程管理的创新之路
人工智能·深度学习·目标检测·机器学习·语言模型·自然语言处理·数据挖掘
小言从不摸鱼4 小时前
【AI大模型】ELMo模型介绍:深度理解语言模型的嵌入艺术
人工智能·深度学习·语言模型·自然语言处理·transformer