删除 pytorch 中张量字典列表中所有张量的某些元素

Question

我是 pytorch 的新手，并尝试通过玩简单的代码来学习。 我所拥有的是一个长度为零的列表，由字典填充，其中该字典中的值是张量。 这是这个列表的样子：

A = [{'boxes': tensor([[ 142.1232,  142.9373, 1106.0452,  971.3792],
    [ 259.1277,  618.4834, 1100.1293, 1028.8989],
    [ 232.1346,  692.5888,  763.3408, 1028.6766],
    [ 206.8070,  312.2080, 1137.1434, 1013.4373],
    [ 495.9471,  675.7287,  978.5932, 1012.7568]], grad_fn=<StackBackward>), 'labels': tensor([16,  1,  1,  1,  1]), 'scores': tensor([0.9988, 0.9489, 0.5228, 0.3500, 0.0639], grad_fn=<IndexBackward>), 'masks': tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]]], grad_fn=<UnsqueezeBackward0>)}]

此列表是 Mask-RCNN 的 output，我想删除嵌套字典中所有张量的某些元素。 在这种情况下，我只想保留与 class label "1" 相关的信息。 与每个 class（“框”、“标签”、“分数”和“掩码”）相关的信息都位于每个张量中的相同位置（索引）。 因此，我找到了带有关键“标签”的张量中所有“1”的索引：

idxOfClass = [i for i, x in enumerate(list(pred[0]['labels'])) if x == 1]

这给了我： [1, 2, 3, 4] 。 然后，我想在嵌套字典中的所有张量中保留位于 idxOfClass 索引处的所有值。 如果我做这样的事情：

Anew = [{pred[0]['boxes'][idxOfClass],pred[0]['labels'][idxOfClass],pred[0]['masks'][idxOfClass],pred[0]['scores'][idxOfClass]}]

我得到：

[{tensor([[ 259.1277,  618.4834, 1100.1293, 1028.8989],
    [ 232.1346,  692.5888,  763.3408, 1028.6766],
    [ 206.8070,  312.2080, 1137.1434, 1013.4373],
    [ 495.9471,  675.7287,  978.5932, 1012.7568]], grad_fn= 
    <IndexBackward>), tensor([0.9489, 0.5228, 0.3500, 0.0639], 
    grad_fn= 
    <IndexBackward>), tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]]], grad_fn=<IndexBackward>), 
    tensor([1, 1, 1, 1])}]

但是，这不是以张量作为字典值的字典列表。 这是张量列表，没有嵌套字典的键值结构。 我的问题是，“当我使用特定元素的索引删除所有张量的某些元素时，有什么方法可以保留列表的原始结构？”。

Answer 1

您可以在构建新的 pred res 时添加键名。

Anew = [{'boxes': pred[0]['boxes'][idxOfClass],'labels': pred[0]['labels'][idxOfClass],'masks': pred[0]['masks'][idxOfClass],'scores': pred[0]['scores'][idxOfClass]}]