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。

相关推荐
惯导马工1 天前
【论文导读】ORB-SLAM3:An Accurate Open-Source Library for Visual, Visual-Inertial and
深度学习·算法
隐语SecretFlow2 天前
国人自研开源隐私计算框架SecretFlow,深度拆解框架及使用【开发者必看】
深度学习
Billy_Zuo2 天前
人工智能深度学习——卷积神经网络(CNN)
人工智能·深度学习·cnn
羊羊小栈2 天前
基于「YOLO目标检测 + 多模态AI分析」的遥感影像目标检测分析系统(vue+flask+数据集+模型训练)
人工智能·深度学习·yolo·目标检测·毕业设计·大作业
l12345sy2 天前
Day24_【深度学习—广播机制】
人工智能·pytorch·深度学习·广播机制
九章云极AladdinEdu2 天前
超参数自动化调优指南:Optuna vs. Ray Tune 对比评测
运维·人工智能·深度学习·ai·自动化·gpu算力
研梦非凡2 天前
ICCV 2025|从粗到细:用于高效3D高斯溅射的可学习离散小波变换
人工智能·深度学习·学习·3d
通街市密人有3 天前
IDF: Iterative Dynamic Filtering Networks for Generalizable Image Denoising
人工智能·深度学习·计算机视觉
智数研析社3 天前
9120 部 TMDb 高分电影数据集 | 7 列全维度指标 (评分 / 热度 / 剧情)+API 权威源 | 电影趋势分析 / 推荐系统 / NLP 建模用
大数据·人工智能·python·深度学习·数据分析·数据集·数据清洗
七元权3 天前
论文阅读-Correlate and Excite
论文阅读·深度学习·注意力机制·双目深度估计