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。

相关推荐
1296004525 小时前
机器学习的可解释性
人工智能·深度学习·自然语言处理·transformer
MYH5166 小时前
神经元激活函数在神经网络里起着关键作用
人工智能·深度学习·神经网络
尘浮7287 小时前
60天python训练计划----day50
人工智能·python·深度学习
强盛小灵通专卖员8 小时前
基于深度学习RT-DETR算法的盲人障碍物目标检测:提升盲人出行安全的智能化突破
深度学习·算法·目标检测·计算机视觉·rt-detr·小论文·计算机期刊
Nelson_hehe9 小时前
论文解析:一文弄懂ResNet(图像识别分类、目标检测)
深度学习·目标检测·计算机视觉·分类·残差网络·resnet
深度之眼9 小时前
贝叶斯深度学习!华科大《Nat. Commun.》发表BNN重大突破!
人工智能·深度学习·神经网络
MYH5169 小时前
深度学习聊天机器人 需要考虑
人工智能·深度学习·机器人
视觉语言导航9 小时前
中南大学无人机智能体的全面评估!BEDI:用于评估无人机上具身智能体的综合性基准测试
人工智能·深度学习·无人机·具身智能
(・Д・)ノ10 小时前
python打卡day47
人工智能·深度学习
不算专业的专业人才10 小时前
AI 驱动全链路后台自动化:从数据破壁到商业闭环的深度解析与实践进阶
深度学习·机器学习·ai编程