含CUDA加载PyTorch张量的Python字典所在设备咨询

1. 字典的存储位置 #

不管是哪种字典，字典本身始终位于CPU内存中，原因如下：

• Python的容器类型（比如字典、列表）属于CPU内存中的数据结构，它们只会持有对其他对象的引用，不会把自身搬到GPU上。
• 两种字典的区别仅在于内部元素的存储位置：
  • 第一种字典：rects、points、scores、image_ids都是PyTorch张量，且这些张量标注了device='cuda:0'，说明这些张量的数据存储在GPU显存中，字典只是保存了这些GPU张量的引用。
    对应字典示例：

    {
        'rects': tensor([[405.5453, 126.3100, 646.7983, 454.6852]], device='cuda:0'),
        'points': tensor([[[477.5181, 247.7193],
                 [591.4567, 247.1899],
                 [544.2579, 328.7857],
                 [494.3871, 371.9922],
                 [583.4272, 370.3429]]], device='cuda:0'),
        'scores': tensor([0.9999], device='cuda:0'),
        'image_ids': tensor([0], device='cuda:0')
    }
    

  • 第二种字典：rects是普通Python列表（存储在CPU内存），points是GPU张量，字典本身仍在CPU内存，同时持有CPU列表和GPU张量的引用。
    对应字典示例：

    {
        'rects': [405.5453, 126.3100, 646.7983, 454.6852],
        'points': tensor([[[477.5181, 247.7193],
                 [591.4567, 247.1899],
                 [544.2579, 328.7857],
                 [494.3871, 371.9922],
                 [583.4272, 370.3429]]], device='cuda:0')
    }
    

2. device="cuda:0"的含义 #

• cuda是NVIDIA推出的GPU计算平台，仅支持NVIDIA系列显卡；
• 数字0是GPU设备的索引：如果机器有多块NVIDIA GPU，会按顺序编号为cuda:0、cuda:1、cuda:2……，单GPU机器上就只有cuda:0；
• 这个属性表示对应的PyTorch张量的数据存储在编号为0的NVIDIA GPU显存中，后续对该张量的运算也会优先在这块GPU上执行。

内容的提问来源于stack exchange，提问作者NJ__